Phase 7

Signed-off-by: Christian Vetter <christian.vetter@here.com>

Author: VeaaC

flatdata-py/flatdata/lib/archive.py

@@ -3,17 +3,26 @@
  See the LICENSE file in the root of this project for license details.
 '''
 
-from collections import namedtuple
-from typing import Any
+from __future__ import annotations
+
+from typing import Any, NamedTuple, TYPE_CHECKING
 
 import pandas as pd
 
 from .errors import MissingResourceError, SchemaMismatchError
 
-ResourceSignature = namedtuple("ResourceSignature",
-                               ["container", "initializer", "schema", "is_optional", "doc"])
+if TYPE_CHECKING:
+    from .resources import ReadStorage, ResourceBase
+
+
+class ResourceSignature(NamedTuple):
+    container: type[ResourceBase] | type[Archive]
+    initializer: Any
+    schema: str
+    is_optional: bool
+    doc: str
 
-def _is_archive_signature(resource_signature: Any) -> bool:
+def _is_archive_signature(resource_signature: ResourceSignature) -> bool:
     return bool(resource_signature.container == Archive)
 
 _SCHEMA_EXT = ".schema"
@@ -26,16 +35,16 @@ class Archive:
     """
     _NAME: str
     _SCHEMA: str
-    _RESOURCES: dict[str, Any]
+    _RESOURCES: dict[str, ResourceSignature]
 
-    def __init__(self, resource_storage: Any) -> None:
+    def __init__(self, resource_storage: ReadStorage) -> None:
         """
         Opens archive from a given resource storage.
         :raises flatdata.errors.CorruptArchiveError
         :raises flatdata.errors.SchemaMismatchError
         :param resource_storage: Resource storage to use.
         """
-        self._resource_storage: Any = resource_storage
+        self._resource_storage: ReadStorage = resource_storage
         self._loaded_resources: dict[str, Any] = {}
 
         # Preload resources and check their schemas
@@ -78,7 +87,7 @@ def resource_schema(cls, resource: str) -> str:
         return str(cls._RESOURCES[resource].schema)
 
     @classmethod
-    def open(cls, storage: Any, name: str, initializer: Any, is_optional: bool = False) -> Any:
+    def open(cls, storage: ReadStorage, name: str, initializer: type[Archive], is_optional: bool = False) -> Archive | None:
         nested_storage = storage.get(name, is_optional)
         assert nested_storage is not None or is_optional
         if nested_storage is None:
@@ -93,7 +102,7 @@ def size_in_bytes(self) -> int:
     def __len__(self) -> int:
         return len(self._RESOURCES)
 
-    def _schema_validated_resource_signature(self, name: str) -> Any:
+    def _schema_validated_resource_signature(self, name: str) -> ResourceSignature | None:
         resource_signature = self._RESOURCES[name]
         # We check only schema for non-subarchives, since the subarchives schema is checked,
         # when it is initialized.
@@ -120,7 +129,7 @@ def _open_resource(self, name: str) -> Any:
         return None
 
     @staticmethod
-    def _check_non_subarchive_schema(name: str, resource_signature: Any, storage: Any) -> None:
+    def _check_non_subarchive_schema(name: str, resource_signature: ResourceSignature, storage: Any) -> None:
         actual_schema = bytes(storage).decode()
         if actual_schema != resource_signature.schema:
             raise SchemaMismatchError(

flatdata-py/flatdata/lib/archive_builder.py

@@ -3,28 +3,43 @@
  See the LICENSE file in the root of this project for license details.
 '''
 
-from collections import namedtuple
+from __future__ import annotations
+
 import os
-from typing import Any
+from typing import Any, NamedTuple, Protocol, TYPE_CHECKING
 
 from .errors import IndexWriterError, MissingFieldError, UnknownFieldError, \
     UnknownStructureError, UnknownResourceError, ResourceAlreadySetError
 
 from .resources import Instance, Vector, Multivector, RawData
 from .data_access import write_value
 
+if TYPE_CHECKING:
+    from .resource_storage import _Resource
+    from .structure import Structure
+
 _SCHEMA_EXT = ".schema"
 
-ResourceSignature = namedtuple("ResourceSignature",
-                               ["container", "initializer", "schema", "is_optional", "doc"])
+
+class ResourceSignature(NamedTuple):
+    container: type
+    initializer: Any
+    schema: str
+    is_optional: bool
+    doc: str
+
+
+class WriteStorage(Protocol):
+    def get(self, resource_name: str, is_subarchive: bool = False) -> Any: ...
+    def close(self) -> None: ...
 
 
 class IndexWriter:
     """
     IndexWriter class. Only applicable when multivector is present in archive schema.
     """
 
-    def __init__(self, name: str, size: int, resource_storage: Any) -> None:
+    def __init__(self, name: str, size: int, resource_storage: WriteStorage) -> None:
         """
         Create IndexWriter class.
 
@@ -37,7 +52,7 @@ def __init__(self, name: str, size: int, resource_storage: Any) -> None:
 
         self._name = name
         self._index_size = size
-        self._fout = resource_storage.get(f'{self._name}_index', False)
+        self._fout: _Resource = resource_storage.get(f'{self._name}_index', False)
 
     def add(self, index: int) -> None:
         """
@@ -63,16 +78,16 @@ class ArchiveBuilder:
     """
     _NAME: str
     _SCHEMA: str
-    _RESOURCES: dict[str, Any]
+    _RESOURCES: dict[str, ResourceSignature]
 
-    def __init__(self, resource_storage: Any, path: str = "") -> None:
+    def __init__(self, resource_storage: WriteStorage, path: str = "") -> None:
         """
         Opens archive from a given resource writer.
         :param resource_storage: storage manager to store and write to disc
         :param path: file path where archive is created
         """
         self._path = os.path.join(path, self._NAME)
-        self._resource_storage = resource_storage
+        self._resource_storage: WriteStorage = resource_storage
         self._write_archive_signature()
         self._write_archive_schema()
         self._resources_written = [f"{self._NAME}.archive"]
@@ -129,7 +144,7 @@ def subarchive(self, name: str) -> 'ArchiveBuilder':
         raise NotImplementedError(f"subarchive '{name}' is not implemented")
 
     @classmethod
-    def __validate_structure_fields(cls, name: str, struct: dict[str, Any], initializer: Any) -> None:
+    def __validate_structure_fields(cls, name: str, struct: dict[str, Any], initializer: type[Structure]) -> None:
         '''
         Validates whether passed object has all required fields
 
@@ -146,7 +161,7 @@ def __validate_structure_fields(cls, name: str, struct: dict[str, Any], initiali
             if key not in initializer._FIELD_KEYS:
                 raise UnknownFieldError(key, name)
 
-    def __set_instance(self, storage: Any, name: str, value: dict[str, Any]) -> None:
+    def __set_instance(self, storage: _Resource, name: str, value: dict[str, Any]) -> None:
         '''
         Creates and writes instance type resource
 
@@ -164,7 +179,7 @@ def __set_instance(self, storage: Any, name: str, value: dict[str, Any]) -> None
 
         storage.write(bout)
 
-    def __set_vector(self, storage: Any, name: str, vector: list[dict[str, Any]]) -> None:
+    def __set_vector(self, storage: _Resource, name: str, vector: list[dict[str, Any]]) -> None:
         '''
         Creates and writes vector resource
 
@@ -183,7 +198,7 @@ def __set_vector(self, storage: Any, name: str, vector: list[dict[str, Any]]) ->
                             field.is_signed, value[key])
             storage.write(bout)
 
-    def __set_multivector(self, storage: Any, name: str, value: list[list[dict[str, Any]]]) -> None:
+    def __set_multivector(self, storage: _Resource, name: str, value: list[list[dict[str, Any]]]) -> None:
         '''
         Creates and writes multivector resource
 

flatdata-py/flatdata/lib/resources.py

@@ -3,23 +3,32 @@
  See the LICENSE file in the root of this project for license details.
 '''
 
+from __future__ import annotations
+
 from collections.abc import Iterator
 import json
-from typing import Any
+from typing import Any, Protocol, TYPE_CHECKING
 
 import pandas as pd
 import numpy as np
 
-from .data_access import read_value, read_field_vectorized
+from .data_access import ReadableBuffer, read_value, read_field_vectorized
 from .errors import CorruptResourceError
 
+if TYPE_CHECKING:
+    from .structure import Structure
+
 SIZE_OFFSET_IN_BITS = 64
 SIZE_OFFSET_IN_BYTES = SIZE_OFFSET_IN_BITS // 8
 SIZE_PADDING_IN_BYTES = 8
 
 
+class ReadStorage(Protocol):
+    def get(self, key: str, is_optional: bool = False) -> Any: ...
+
+
 class ResourceBase:
-    def __init__(self, mem: Any, element_type: type[Any]) -> None:
+    def __init__(self, mem: ReadableBuffer, element_type: type[Structure] | None) -> None:
         if len(mem) < (SIZE_OFFSET_IN_BYTES + SIZE_PADDING_IN_BYTES):
             raise CorruptResourceError()
         self._mem = memoryview(mem)
@@ -32,9 +41,11 @@ def size_in_bytes(self) -> int:
         return len(self._mem)
 
     def _item_offset(self, index: int) -> int:
+        assert self._element_type is not None
         return int(SIZE_OFFSET_IN_BYTES + self._element_type._SIZE_IN_BYTES * index)
 
     def _get_item(self, index: int) -> Any:
+        assert self._element_type is not None
         offset = self._item_offset(index)
         return self._element_type(self._mem, offset)
 
@@ -67,7 +78,7 @@ def __repr__(self) -> str:
         return json.dumps(self._repr_attributes(), indent=4)
 
     @classmethod
-    def open(cls, storage: Any, name: str, initializer: Any, is_optional: bool = False) -> Any:
+    def open(cls, storage: ReadStorage, name: str, initializer: Any, is_optional: bool = False) -> Any:
         return cls(storage.get(name, is_optional), initializer)
 
 
@@ -82,6 +93,7 @@ def to_numpy(self, limit: int | None = None) -> Any:
         if limit is not None:
             sliced = sliced[:limit]
 
+        assert self._sequence._element_type is not None
         fields = self._sequence._element_type._FIELDS
         dtype = self._sequence._element_type.dtype()
         result = np.empty(sliced.shape[0], dtype=dtype)
@@ -99,6 +111,7 @@ def __iter__(self) -> Iterator[Any]:
             yield self._sequence[i]
 
     def __getattr__(self, name: str) -> pd.DataFrame:
+        assert self._sequence._element_type is not None
         try:
             field = self._sequence._element_type._FIELDS[name]
         except KeyError:
@@ -112,7 +125,7 @@ def __repr__(self) -> str:
 
 
 class Vector(ResourceBase):
-    def __init__(self, mem: Any, element_type: type[Any]) -> None:
+    def __init__(self, mem: ReadableBuffer, element_type: type[Structure]) -> None:
         ResourceBase.__init__(self, mem, element_type)
         size_in_bytes = read_value(self._mem, 0, SIZE_OFFSET_IN_BITS, False)
         size, rem = divmod(size_in_bytes, self._type_size_in_bytes)
@@ -122,6 +135,7 @@ def __init__(self, mem: Any, element_type: type[Any]) -> None:
     def to_numpy(self) -> Any:
         """Convert entire vector to a numpy structured array (vectorized)."""
         raw_2d = self._as_numpy_2d()
+        assert self._element_type is not None
         fields = self._element_type._FIELDS
         dtype = self._element_type.dtype()
         result = np.empty(self._size, dtype=dtype)
@@ -147,11 +161,13 @@ def __getitem__(self, index: int | slice) -> Any:
     def __iter__(self) -> Iterator[Any]:
         mem = self._mem
         element_type = self._element_type
+        assert element_type is not None
         size_bytes = self._type_size_in_bytes
         for i in range(self._size):
             yield element_type(mem, SIZE_OFFSET_IN_BYTES + size_bytes * i)
 
     def __getattr__(self, name: str) -> pd.DataFrame:
+        assert self._element_type is not None
         try:
             field = self._element_type._FIELDS[name]
         except KeyError:
@@ -178,14 +194,14 @@ def __repr__(self) -> str:
 
 
 class Multivector(ResourceBase):
-    def __init__(self, index_mem: Any, mem: Any, index_type: type[Any], *element_types: type[Any]) -> None:
+    def __init__(self, index_mem: ReadableBuffer, mem: ReadableBuffer, index_type: type[Structure], *element_types: type[Structure]) -> None:
         self._index = Vector(index_mem, index_type)
-        self._mem = mem
+        self._mem = memoryview(mem)
         self._element_types = list(element_types)
         self._index_type = index_type
 
     @classmethod
-    def open(cls, storage: Any, name: str, initializer: Any, is_optional: bool = False) -> 'Multivector':
+    def open(cls, storage: ReadStorage, name: str, initializer: list[type[Structure]], is_optional: bool = False) -> Multivector:
         return cls(storage.get(name + "_index", is_optional),
                    storage.get(name, is_optional),
                    *initializer)
@@ -208,6 +224,7 @@ def __getitem__(self, index: int | slice) -> Any:
             type_index = read_value(self._mem, offset * 8, 8, False)
             offset += 1
             element_type = self._element_types[type_index]
+            assert element_type is not None
             element = element_type(self._mem, offset)
             elements.append(element)
             offset += element_type._SIZE_IN_BYTES
@@ -280,6 +297,7 @@ def __iter__(self) -> Iterator[Any]:
             yield self._get_item(i)
 
     def __getattr__(self, name: str) -> Any:
+        assert self._element_type is not None
         offset = self._item_offset(0)
         return getattr(self._element_type(self._mem, offset), name)
 

flatdata-py/flatdata/lib/structure.py

@@ -17,6 +17,8 @@ class Structure:
     _FIELDS: dict[str, FieldSignature]
     _FIELD_KEYS: list[str]
     _SCHEMA: str
+    _SIZE_IN_BYTES: int
+    _NAME: str
 
     def __init_subclass__(cls, **kwargs: Any) -> None:
         super().__init_subclass__(**kwargs)