fix: allow nested db layer entries

ZettaAI · Feb 24, 2025 · 0e439a9 · 0e439a9
1 parent e6fc726
commit 0e439a9
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 7 deletions.
diff --git a/tests/unit/layer/db_layer/firestore/test_backend.py b/tests/unit/layer/db_layer/firestore/test_backend.py
@@ -140,6 +140,11 @@ def test_batching(firestore_emulator, mocker) -> None:
 
     ROW_COUNT = 150
     COLS = ("col_a", "col_b")
+    mocker.patch(
+        "zetta_utils.layer.db_layer.firestore.backend.FirestoreBackend.__len__",
+        return_value=0,
+    )
+    assert layer.get_batch(10, 10) == {}
     mocker.patch(
         "zetta_utils.layer.db_layer.firestore.backend.FirestoreBackend.__len__",
         return_value=ROW_COUNT,

diff --git a/zetta_utils/layer/db_layer/__init__.py b/zetta_utils/layer/db_layer/__init__.py
@@ -1,5 +1,5 @@
 from .index import DBIndex
-from .backend import DBDataT, DBBackend, DBArrayValueT, DBValueT, DBRowDataT
+from .backend import DBDataT, DBBackend, DBValueT, DBRowDataT
 
 from .layer import DBLayer, UserDBIndex, ColIndex, DBDataProcT
 

diff --git a/zetta_utils/layer/db_layer/backend.py b/zetta_utils/layer/db_layer/backend.py
@@ -2,14 +2,12 @@
 from __future__ import annotations
 
 from abc import abstractmethod
-from typing import MutableMapping, Sequence, Union
+from typing import Mapping, MutableMapping, Sequence, Union
 
 from .. import Backend
 from . import DBIndex
 
-DBScalarValueT = Union[bool, int, float, str]
-DBArrayValueT = list[DBScalarValueT]
-DBValueT = Union[DBScalarValueT, DBArrayValueT]
+DBValueT = Union[bool, int, float, str, Sequence["DBValueT"], Mapping[str, "DBValueT"]]
 DBRowDataT = MutableMapping[str, DBValueT | None]
 DBDataT = Sequence[DBRowDataT]
 

diff --git a/zetta_utils/layer/db_layer/firestore/backend.py b/zetta_utils/layer/db_layer/firestore/backend.py
@@ -16,7 +16,6 @@
     stop_after_attempt,
     wait_exponential,
 )
-from typeguard import typechecked
 
 from zetta_utils import builder
 from zetta_utils.layer.db_layer import DBBackend, DBDataT, DBIndex, DBRowDataT
@@ -26,7 +25,6 @@
 
 
 @builder.register("FirestoreBackend")
-@typechecked
 @attrs.mutable
 class FirestoreBackend(DBBackend):
     """