Skip to content

Commit

Permalink
add flatten method to collection options structures
Browse files Browse the repository at this point in the history
  • Loading branch information
@hemidactylus
hemidactylus committed Mar 27, 2024
1 parent a185212 commit 5cd4ae7
Show file tree
Hide file tree
Showing 2 changed files with 180 additions and 62 deletions.
52 changes: 52 additions & 0 deletions astrapy/info.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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:
Expand All @@ -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 {
Expand All @@ -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"),
Expand Down Expand Up @@ -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 {
Expand All @@ -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"),
Expand Down Expand Up @@ -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 {
Expand All @@ -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 {}),
Expand Down
190 changes: 128 additions & 62 deletions tests/idiomatic/unit/test_collection_options.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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

0 comments on commit 5cd4ae7

Please sign in to comment.