feat: typed annotations

tests/unit/db_annotations/conftest.py

@@ -17,18 +17,25 @@ def annotations_db(firestore_emulator):
 def collections_db(firestore_emulator):
     db = build_firestore_layer(collection.DB_NAME, project=firestore_emulator)
     collection.COLLECTIONS_DB = db
+    collections = collection.read_collections()
+    collection.delete_collections(list(collections.keys()))
     return collection.COLLECTIONS_DB
 
 
 @pytest.fixture(scope="session")
 def layer_groups_db(firestore_emulator):
     db = build_firestore_layer(layer_group.DB_NAME, project=firestore_emulator)
     layer_group.LAYER_GROUPS_DB = db
+    layer_groups = layer_group.read_layer_groups()
+    layer_group.delete_layer_groups(list(layer_groups.keys()))
     return layer_group.LAYER_GROUPS_DB
 
 
 @pytest.fixture(scope="session")
 def layers_db(firestore_emulator):
     db = build_firestore_layer(layer.DB_NAME, project=firestore_emulator)
     layer.LAYERS_DB = db
+    layers = layer.read_layers()
+    for e in layers:
+        layer.delete_layer(e)
     return layer.LAYERS_DB

tests/unit/db_annotations/test_annotation.py

@@ -1,4 +1,4 @@
-# pylint: disable=unused-argument
+# pylint: disable=unused-argument,redefined-outer-name
 
 from typing import cast
 
@@ -7,31 +7,43 @@
 from zetta_utils.db_annotations import annotation, collection, layer, layer_group
 
 
-def _init_collection_and_layer_group(
-    collection_name: str, layer_group_name: str
-) -> tuple[str, str]:
+@pytest.fixture
+def collection_and_layer_group():
     user = "john_doe"
-    collection_id = collection.add_collection(collection_name, user, "this is a test")
+    collection_id = collection.add_collection("collection_x0", user, "this is a test")
+    layer_groups = layer_group.read_layer_groups(collection_ids=[collection_id])
+    annotations = annotation.read_annotations(collection_ids=[collection_id])
+    layer_group.delete_layer_groups(list(layer_groups.keys()))
+    annotation.delete_annotations(list(annotations.keys()))
 
     layer_id0 = layer.add_layer("test_layer0", "precomputed://test0", "this is a test")
     layer_id1 = layer.add_layer("test_layer1", "precomputed://test1", "this is a test")
 
     layer_group_id = layer_group.add_layer_group(
-        name=layer_group_name,
+        name="layer_group_x0",
         collection_id=collection_id,
         user=user,
         layers=[layer_id0, layer_id1],
         comment="this is a test",
     )
-    return (collection_id, layer_group_id)
+    yield (collection_id, layer_group_id)
+    layer_group.delete_layer_group(layer_group_id)
+    layer_groups = layer_group.read_layer_groups(collection_ids=[collection_id])
+    annotations = annotation.read_annotations(collection_ids=[collection_id])
+    layer_group.delete_layer_groups(list(layer_groups.keys()))
+    annotation.delete_annotations(list(annotations.keys()))
+    collection.delete_collection(collection_id=collection_id)
 
 
 def test_add_update_delete_annotation(
-    firestore_emulator, annotations_db, collections_db, layer_groups_db, layers_db
+    firestore_emulator,
+    annotations_db,
+    collections_db,
+    layer_groups_db,
+    layers_db,
+    collection_and_layer_group,
 ):
-    collection_id, layer_group_id = _init_collection_and_layer_group(
-        "new_collection0", "new_lgroup0"
-    )
+    collection_id, layer_group_id = collection_and_layer_group
     annotation_raw = {
         "pointA": [1, 1, 1],
         "pointB": [1, 1, 5],
@@ -47,9 +59,9 @@ def test_add_update_delete_annotation(
     )
 
     _annotation = annotation.read_annotation(_id)
-    assert _annotation["collection"] == collection_id
-    assert _annotation["layer_group"] == layer_group_id
-    assert len(cast(list, _annotation["tags"])) == 2
+    assert _annotation.collection_id == collection_id
+    assert _annotation.layer_group_id == layer_group_id
+    assert len(cast(list, _annotation.tags)) == 2
 
     annotation.update_annotation(
         _id,
@@ -59,20 +71,23 @@ def test_add_update_delete_annotation(
         tags=["tag2"],
     )
     _annotation = annotation.read_annotation(_id)
-    assert len(cast(list, _annotation["tags"])) == 1
-    assert cast(list, _annotation["tags"])[0] == "tag2"
+    assert len(cast(list, _annotation.tags)) == 1
+    assert cast(list, _annotation.tags)[0] == "tag2"
 
     annotation.delete_annotation(_id)
     with pytest.raises(KeyError):
         annotation.read_annotation(_id)
 
 
 def test_add_update_annotations(
-    firestore_emulator, annotations_db, collections_db, layer_groups_db, layers_db
+    firestore_emulator,
+    annotations_db,
+    collections_db,
+    layer_groups_db,
+    layers_db,
+    collection_and_layer_group,
 ):
-    collection_id, layer_group_id = _init_collection_and_layer_group(
-        "new_collection1", "new_lgroup1"
-    )
+    collection_id, layer_group_id = collection_and_layer_group
     annotations_raw = [
         {
             "pointA": [1, 1, 1],
@@ -107,10 +122,10 @@ def test_add_update_annotations(
     )
 
     _annotations = annotation.read_annotations(annotation_ids=_ids)
-    assert _annotations[0]["type"] == "line"
-    assert _annotations[1]["type"] == "point"
-    assert len(cast(list, _annotations[0]["tags"])) == 2
-    assert len(cast(list, _annotations[1]["tags"])) == 2
+    assert _annotations[0].ng_annotation.type == "line"
+    assert _annotations[1].ng_annotation.type == "point"
+    assert len(cast(list, _annotations[0].tags)) == 2
+    assert len(cast(list, _annotations[1].tags)) == 2
 
     annotation.update_annotations(
         _ids,
@@ -120,14 +135,12 @@ def test_add_update_annotations(
         tags=["tag2"],
     )
     _annotations = annotation.read_annotations(annotation_ids=_ids)
-    assert len(cast(list, _annotations[0]["tags"])) == 1
-    assert len(cast(list, _annotations[1]["tags"])) == 1
+    assert len(cast(list, _annotations[0].tags)) == 1
+    assert len(cast(list, _annotations[1].tags)) == 1
 
 
-def test_read_delete_annotations(firestore_emulator, annotations_db):
-    collection_id, layer_group_id = _init_collection_and_layer_group(
-        "new_collection2", "new_lgroup2"
-    )
+def test_read_delete_annotations(firestore_emulator, annotations_db, collection_and_layer_group):
+    collection_id, layer_group_id = collection_and_layer_group
     ng_annotations = annotation.parse_ng_annotations(
         [
             {

tests/unit/db_annotations/test_collection.py

@@ -2,7 +2,34 @@
 
 import pytest
 
-from zetta_utils.db_annotations import collection
+from zetta_utils.db_annotations import annotation, collection, layer_group
+
+
+@pytest.fixture
+def collection_x0():
+    yield from collection_maker("collection_x0")
+
+
+@pytest.fixture
+def collection_x1():
+    yield from collection_maker("collection_x1")
+
+
+def collection_maker(collection_name: str):
+    user = "john_doe"
+    collection_id = collection.add_collection(collection_name, user, "this is a test")
+    layer_groups = layer_group.read_layer_groups(collection_ids=[collection_id])
+    annotations = annotation.read_annotations(collection_ids=[collection_id])
+    layer_group.delete_layer_groups(list(layer_groups.keys()))
+    annotation.delete_annotations(list(annotations.keys()))
+
+    yield collection_id
+
+    layer_groups = layer_group.read_layer_groups(collection_ids=[collection_id])
+    annotations = annotation.read_annotations(collection_ids=[collection_id])
+    layer_group.delete_layer_groups(list(layer_groups.keys()))
+    annotation.delete_annotations(list(annotations.keys()))
+    collection.delete_collection(collection_id=collection_id)
 
 
 def test_add_update_delete_collection(firestore_emulator, collections_db):
@@ -14,14 +41,16 @@ def test_add_update_delete_collection(firestore_emulator, collections_db):
         collection.add_collection(old_name, old_user, "this is a test")
 
     _collection = collection.read_collection(_id)
-    assert _collection["name"] == old_name
-    assert _collection["created_by"] == old_user
+    assert _collection.name == old_name
+    assert _collection.created_by == old_user
 
     new_user = "jane_doe"
-    collection.update_collection(_id, new_user, "test_collection1", comment="this is also a test")
+    collection.update_collection(
+        _id, user=new_user, name="test_collection1", comment="this is also a test"
+    )
     _collection = collection.read_collection(_id)
-    assert _collection["name"] != old_name
-    assert _collection["modified_by"] == new_user
+    assert _collection.name != old_name
+    assert _collection.modified_by == new_user
 
     collection.delete_collection(_id)
     with pytest.raises(KeyError):
@@ -33,13 +62,13 @@ def test_read_delete_collections(firestore_emulator, collections_db):
     collection_id1 = collection.add_collection("test_collection1", "jane_doe", "this is a test")
     _collections = collection.read_collections(collection_ids=[collection_id0, collection_id1])
 
-    assert _collections[0]["name"] == "test_collection0"
-    assert _collections[0]["created_by"] == "john_doe"
-    assert _collections[1]["name"] == "test_collection1"
-    assert _collections[1]["created_by"] == "jane_doe"
+    assert _collections[0].name == "test_collection0"
+    assert _collections[0].created_by == "john_doe"
+    assert _collections[1].name == "test_collection1"
+    assert _collections[1].created_by == "jane_doe"
 
     _collections1 = collection.read_collections()
-    assert len(_collections1) == 5
+    assert len(_collections1) == 2
 
     collection.delete_collections([collection_id0, collection_id1])
 

tests/unit/db_annotations/test_layer.py

@@ -7,25 +7,25 @@ def test_add_update_layer(firestore_emulator, layers_db):
     old_name = "test_layer0"
     _id = layer.add_layer(old_name, "precomputed://test", comment="this is a test")
     _layer = layer.read_layer(_id)
-    assert _layer["name"] == old_name
-    assert _layer["source"] == "precomputed://test"
+    assert _layer.name == old_name
+    assert _layer.source == "precomputed://test"
 
     layer.update_layer(
         _id, name="test_layer1", source="precomputed://test", comment="this is a test"
     )
     _layer = layer.read_layer(_id)
-    assert _layer["name"] != old_name
+    assert _layer.name != old_name
 
 
 def test_read_layers(firestore_emulator, layers_db):
     layer_id0 = layer.add_layer("test_layer0", "precomputed://test0", "this is a test")
     layer_id1 = layer.add_layer("test_layer1", "precomputed://test1", "this is a test")
     _layers = layer.read_layers(layer_ids=[layer_id0, layer_id1])
 
-    assert _layers[0]["name"] == "test_layer0"
-    assert _layers[0]["source"] == "precomputed://test0"
-    assert _layers[1]["name"] == "test_layer1"
-    assert _layers[1]["source"] == "precomputed://test1"
+    assert _layers[0].name == "test_layer0"
+    assert _layers[0].source == "precomputed://test0"
+    assert _layers[1].name == "test_layer1"
+    assert _layers[1].source == "precomputed://test1"
 
     _layers1 = layer.read_layers()
     assert len(_layers1) == 9

tests/unit/db_annotations/test_layer_group.py

@@ -35,8 +35,8 @@ def test_add_update_delete_layer_group(firestore_emulator, layer_groups_db):
         )
 
     _layer_group = layer_group.read_layer_group(_id)
-    assert _layer_group["name"] == old_name
-    assert len(cast(list, _layer_group["layers"])) == 2
+    assert _layer_group.name == old_name
+    assert len(cast(list, _layer_group.layers)) == 2
 
     layer_group.update_layer_group(
         _id,
@@ -47,7 +47,7 @@ def test_add_update_delete_layer_group(firestore_emulator, layer_groups_db):
         comment="this is a test",
     )
     _layer = layer_group.read_layer_group(_id)
-    assert _layer["name"] != old_name
+    assert _layer.name != old_name
 
     layer_group.delete_layer_group(_id)
     with pytest.raises(KeyError):
@@ -80,13 +80,11 @@ def test_read_delete_layer_groups(firestore_emulator, layer_groups_db):
     )
 
     _layer_groups = layer_group.read_layer_groups(layer_group_ids=[_id0, _id1])
-    assert _layer_groups[0]["name"] == "test_layer_group0"
-    assert len(cast(list, _layer_groups[0]["layers"])) == len(
-        cast(list, _layer_groups[1]["layers"])
-    )
+    assert _layer_groups[0].name == "test_layer_group0"
+    assert len(cast(list, _layer_groups[0].layers)) == len(cast(list, _layer_groups[1].layers))
 
     _layer_groups1 = layer_group.read_layer_groups()
-    assert len(_layer_groups1) == 5
+    assert len(_layer_groups1) == 2
 
     _layer_groups2 = layer_group.read_layer_groups(collection_ids=[collection_id])
     assert len(_layer_groups2) == 2

web_api/app/annotations.py

@@ -83,7 +83,7 @@ async def read_multiple(
     )
     response = []
     for _id, annotation in annotations.items():
-        annotation["id"] = _id
+        annotation.id = _id
         response.append(annotation)
     return response
 

web_api/app/collections.py

@@ -61,7 +61,7 @@ async def read_multiple(collection_ids: Annotated[list[str] | None, Query()] = N
     collections = read_collections()
     response = []
     for _id, collection in collections.items():
-        collection["id"] = _id
+        collection.id = _id
         response.append(collection)
     return response
 

web_api/app/layer_groups.py

@@ -82,7 +82,7 @@ async def read_multiple(
         layer_groups = read_layer_groups()
     response = []
     for _id, layer_group in layer_groups.items():
-        layer_group["id"] = _id
+        layer_group.id = _id
         response.append(layer_group)
     return response
 

web_api/app/layers.py

@@ -44,6 +44,6 @@ async def read_multiple(layer_ids: Annotated[list[str] | None, Query()] = None):
     layers = read_layers()
     response = []
     for _id, layer in layers.items():
-        layer["id"] = _id
+        layer.id = _id
         response.append(layer)
     return response

zetta_utils/db_annotations/__init__.py

@@ -0,0 +1,13 @@
+from .annotation import (
+    read_annotation,
+    read_annotations,
+    add_annotation,
+    add_annotations,
+    update_annotation,
+    update_annotations,
+    parse_ng_annotations,
+)
+from .collection import read_collection, read_collections, add_collection, update_collection
+from .layer_group import read_layer_group, read_layer_groups, add_layer_group, update_layer_group
+from .layer import read_layer, read_layers, add_layer, update_layer
+from . import precomp_annotations