diff --git a/astrapy/info.py b/astrapy/info.py index 0afc11b2..8f46998e 100644 --- a/astrapy/info.py +++ b/astrapy/info.py @@ -142,6 +142,7 @@ class CollectionDefaultIDOptions: def as_dict(self) -> Dict[str, Any]: """Recast this object into a dictionary.""" + return {"type": self.default_id_type} @staticmethod @@ -152,6 +153,7 @@ def from_dict( Create an instance of CollectionDefaultIDOptions from a dictionary such as one from the Data API. """ + if raw_dict is not None: return CollectionDefaultIDOptions(default_id_type=raw_dict["type"]) else: @@ -175,6 +177,7 @@ class CollectionVectorOptions: def as_dict(self) -> Dict[str, Any]: """Recast this object into a dictionary.""" + return { k: v for k, v in { @@ -192,6 +195,7 @@ def from_dict( Create an instance of CollectionVectorOptions from a dictionary such as one from the Data API. """ + if raw_dict is not None: return CollectionVectorOptions( dimension=raw_dict.get("dimension"), @@ -241,6 +245,7 @@ def __repr__(self) -> str: def as_dict(self) -> Dict[str, Any]: """Recast this object into a dictionary.""" + return { k: v for k, v in { @@ -253,12 +258,46 @@ def as_dict(self) -> Dict[str, Any]: if v is not None } + def flatten(self) -> Dict[str, Any]: + """ + Recast this object as a flat key-value pair suitable for + use as kwargs in a create_collection method call. + """ + + _dimension: Optional[int] + _metric: Optional[str] + _indexing: Optional[Dict[str, Any]] + _default_id_type: Optional[str] + if self.vector is not None: + _dimension = self.vector.dimension + _metric = self.vector.metric + else: + _dimension = None + _metric = None + _indexing = self.indexing + if self.default_id is not None: + _default_id_type = self.default_id.default_id_type + else: + _default_id_type = None + + return { + k: v + for k, v in { + "dimension": _dimension, + "metric": _metric, + "indexing": _indexing, + "default_id_type": _default_id_type, + }.items() + if v is not None + } + @staticmethod def from_dict(raw_dict: Dict[str, Any]) -> CollectionOptions: """ Create an instance of CollectionOptions from a dictionary such as one from the Data API. """ + return CollectionOptions( vector=CollectionVectorOptions.from_dict(raw_dict.get("vector")), indexing=raw_dict.get("indexing"), @@ -295,6 +334,7 @@ def as_dict(self) -> Dict[str, Any]: Recast this object into a dictionary. Empty `options` will not be returned at all. """ + return { k: v for k, v in { @@ -304,12 +344,24 @@ def as_dict(self) -> Dict[str, Any]: if v } + def flatten(self) -> Dict[str, Any]: + """ + Recast this object as a flat key-value pair suitable for + use as kwargs in a create_collection method call. + """ + + return { + **(self.options.flatten()), + **{"name": self.name}, + } + @staticmethod def from_dict(raw_dict: Dict[str, Any]) -> CollectionDescriptor: """ Create an instance of CollectionDescriptor from a dictionary such as one from the Data API. """ + return CollectionDescriptor( name=raw_dict["name"], options=CollectionOptions.from_dict(raw_dict.get("options") or {}), diff --git a/tests/idiomatic/unit/test_collection_options.py b/tests/idiomatic/unit/test_collection_options.py index cd86436c..140ef76a 100644 --- a/tests/idiomatic/unit/test_collection_options.py +++ b/tests/idiomatic/unit/test_collection_options.py @@ -16,7 +16,7 @@ Unit tests for the validation/parsing of collection options """ -from typing import Any, Dict, List +from typing import Any, Dict, List, Tuple import pytest @@ -25,91 +25,157 @@ @pytest.mark.describe("test of recasting the collection options from the api") def test_recast_api_collection_dict() -> None: - api_coll_descs: List[Dict[str, Any]] = [ + api_coll_descs: List[Tuple[Dict[str, Any], Dict[str, Any]]] = [ # minimal: - { - "name": "dvv", - }, + ( + { + "name": "col_name", + }, + {"name": "col_name"}, + ), # full: - { - "name": "dvv", - "options": { - "vector": { - "dimension": 1024, - "metric": "cosine", + ( + { + "name": "col_name", + "options": { + "vector": { + "dimension": 1024, + "metric": "cosine", + }, + "indexing": {"deny": ["a"]}, + "defaultId": {"type": "objectId"}, }, + }, + { + "name": "col_name", + "dimension": 1024, + "metric": "cosine", "indexing": {"deny": ["a"]}, - "defaultId": {"type": "objectId"}, + "default_id_type": "objectId", }, - }, + ), # partial/absent 'vector': - { - "name": "dvv", - "options": { - "vector": { - "metric": "cosine", + ( + { + "name": "col_name", + "options": { + "vector": { + "metric": "cosine", + }, + "indexing": {"deny": ["a"]}, + "defaultId": {"type": "objectId"}, }, + }, + { + "name": "col_name", + "metric": "cosine", "indexing": {"deny": ["a"]}, - "defaultId": {"type": "objectId"}, - }, - }, - { - "name": "dvv", - "options": { - "vector": { - "dimension": 1024, + "default_id_type": "objectId", + }, + ), + ( + { + "name": "col_name", + "options": { + "vector": { + "dimension": 1024, + }, + "indexing": {"deny": ["a"]}, + "defaultId": {"type": "objectId"}, }, + }, + { + "name": "col_name", + "dimension": 1024, "indexing": {"deny": ["a"]}, - "defaultId": {"type": "objectId"}, + "default_id_type": "objectId", + }, + ), + ( + { + "name": "col_name", + "options": { + "vector": {}, + "indexing": {"deny": ["a"]}, + "defaultId": {"type": "objectId"}, + }, }, - }, - { - "name": "dvv", - "options": { - "vector": {}, + { + "name": "col_name", "indexing": {"deny": ["a"]}, - "defaultId": {"type": "objectId"}, + "default_id_type": "objectId", + }, + ), + ( + { + "name": "col_name", + "options": { + "indexing": {"deny": ["a"]}, + "defaultId": {"type": "objectId"}, + }, }, - }, - { - "name": "dvv", - "options": { + { + "name": "col_name", "indexing": {"deny": ["a"]}, - "defaultId": {"type": "objectId"}, + "default_id_type": "objectId", }, - }, + ), # no indexing: - { - "name": "dvv", - "options": { - "vector": { - "dimension": 1024, - "metric": "cosine", + ( + { + "name": "col_name", + "options": { + "vector": { + "dimension": 1024, + "metric": "cosine", + }, + "defaultId": {"type": "objectId"}, }, - "defaultId": {"type": "objectId"}, }, - }, + { + "name": "col_name", + "dimension": 1024, + "metric": "cosine", + "default_id_type": "objectId", + }, + ), # no defaultId: - { - "name": "dvv", - "options": { - "vector": { - "dimension": 1024, - "metric": "cosine", + ( + { + "name": "col_name", + "options": { + "vector": { + "dimension": 1024, + "metric": "cosine", + }, + "indexing": {"deny": ["a"]}, }, + }, + { + "name": "col_name", + "dimension": 1024, + "metric": "cosine", "indexing": {"deny": ["a"]}, }, - }, + ), # no indexing + no defaultId: - { - "name": "dvv", - "options": { - "vector": { - "dimension": 1024, - "metric": "cosine", + ( + { + "name": "col_name", + "options": { + "vector": { + "dimension": 1024, + "metric": "cosine", + }, }, }, - }, + { + "name": "col_name", + "dimension": 1024, + "metric": "cosine", + }, + ), ] - for api_coll_desc in api_coll_descs: + for api_coll_desc, flattened_dict in api_coll_descs: assert CollectionDescriptor.from_dict(api_coll_desc).as_dict() == api_coll_desc + assert CollectionDescriptor.from_dict(api_coll_desc).flatten() == flattened_dict