Phase 2

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

Author: VeaaC

flatdata-py/flatdata/lib/archive.py

@@ -13,8 +13,8 @@
 ResourceSignature = namedtuple("ResourceSignature",
                                ["container", "initializer", "schema", "is_optional", "doc"])
 
-def _is_archive_signature(resource_signature):
-    return resource_signature.container == Archive
+def _is_archive_signature(resource_signature: Any) -> bool:
+    return bool(resource_signature.container == Archive)
 
 _SCHEMA_EXT = ".schema"
 
@@ -28,34 +28,34 @@ class Archive:
     _SCHEMA: str
     _RESOURCES: dict[str, Any]
 
-    def __init__(self, resource_storage):
+    def __init__(self, resource_storage: Any) -> 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 = resource_storage
-        self._loaded_resources = {}
+        self._resource_storage: Any = resource_storage
+        self._loaded_resources: dict[str, Any] = {}
 
         # Preload resources and check their schemas
         for name, _ in sorted(list(self._RESOURCES.items())):
             self.__getattr__(name)
 
-    def __getattr__(self, name):
+    def __getattr__(self, name: str) -> Any:
         if name not in self._RESOURCES:
             raise AttributeError("Resource %s not defined in archive." % name)
         if name not in self._loaded_resources:
             self._loaded_resources[name] = self._open_resource(name)
         return self._loaded_resources[name]
 
-    def __dir__(self):
+    def __dir__(self) -> list[str]:
         return list(self._RESOURCES.keys()) + ['schema']
 
-    def __repr__(self):
-        return self.to_data_frame().__repr__()
+    def __repr__(self) -> str:
+        return repr(self.to_data_frame())
 
-    def to_data_frame(self):
+    def to_data_frame(self) -> pd.DataFrame:
         result = []
         for name, signature in self._RESOURCES.items():
             resource = self.__getattr__(name)
@@ -66,34 +66,34 @@ def to_data_frame(self):
                             columns=["Name", "Type", "Optional", "SizeInBytes", "Size"])
 
     @classmethod
-    def name(cls):
+    def name(cls) -> str:
         return cls._NAME
 
     @classmethod
-    def schema(cls):
+    def schema(cls) -> str:
         return cls._SCHEMA
 
     @classmethod
-    def resource_schema(cls, resource):
-        return cls._RESOURCES[resource].schema
+    def resource_schema(cls, resource: str) -> str:
+        return str(cls._RESOURCES[resource].schema)
 
     @classmethod
-    def open(cls, storage, name, initializer, is_optional=False):
+    def open(cls, storage: Any, name: str, initializer: Any, is_optional: bool = False) -> Any:
         nested_storage = storage.get(name, is_optional)
         assert nested_storage is not None or is_optional
         if nested_storage is None:
             return None
         return initializer(nested_storage)
 
-    def size_in_bytes(self):
+    def size_in_bytes(self) -> int:
         return sum(resource_value.size_in_bytes() for resource_value in
                    (self.__getattr__(resource) for resource in self._RESOURCES.keys())
                    if resource_value)
 
-    def __len__(self):
+    def __len__(self) -> int:
         return len(self._RESOURCES)
 
-    def _schema_validated_resource_signature(self, name):
+    def _schema_validated_resource_signature(self, name: str) -> Any:
         resource_signature = self._RESOURCES[name]
         # We check only schema for non-subarchives, since the subarchives schema is checked,
         # when it is initialized.
@@ -107,7 +107,7 @@ def _schema_validated_resource_signature(self, name):
                 return None
         return resource_signature
 
-    def _open_resource(self, name):
+    def _open_resource(self, name: str) -> Any:
         resource_signature = self._schema_validated_resource_signature(name)
         if resource_signature:
             resource = resource_signature.container.open(storage=self._resource_storage,
@@ -120,7 +120,7 @@ def _open_resource(self, name):
         return None
 
     @staticmethod
-    def _check_non_subarchive_schema(name, resource_signature, storage):
+    def _check_non_subarchive_schema(name: str, resource_signature: Any, storage: Any) -> None:
         actual_schema = bytes(storage).decode()
         if actual_schema != resource_signature.schema:
             raise SchemaMismatchError(

flatdata-py/flatdata/lib/archive_builder.py

@@ -24,7 +24,7 @@ class IndexWriter:
     IndexWriter class. Only applicable when multivector is present in archive schema.
     """
 
-    def __init__(self, name, size, resource_storage):
+    def __init__(self, name: str, size: int, resource_storage: Any) -> None:
         """
         Create IndexWriter class.
 
@@ -39,15 +39,15 @@ def __init__(self, name, size, resource_storage):
         self._index_size = size
         self._fout = resource_storage.get(f'{self._name}_index', False)
 
-    def add(self, index):
+    def add(self, index: int) -> None:
         """
         Convert index(number) to bytearray and add to in memory store
         """
         index_bytes = int(index).to_bytes(self._index_size,
                                           byteorder="little", signed=False)
         self._fout.write(index_bytes)
 
-    def finish(self):
+    def finish(self) -> None:
         """
         Complete index resource by adding size and padding followed by writing to file
         """
@@ -65,7 +65,7 @@ class ArchiveBuilder:
     _SCHEMA: str
     _RESOURCES: dict[str, Any]
 
-    def __init__(self, resource_storage, path=""):
+    def __init__(self, resource_storage: Any, path: str = "") -> None:
         """
         Opens archive from a given resource writer.
         :param resource_storage: storage manager to store and write to disc
@@ -78,16 +78,16 @@ def __init__(self, resource_storage, path=""):
         self._resources_written = [f"{self._NAME}.archive"]
 
     @classmethod
-    def name(cls):
+    def name(cls) -> str:
         '''Returns archive name'''
         return cls._NAME
 
     @classmethod
-    def schema(cls):
+    def schema(cls) -> str:
         '''Returns archive schema'''
         return cls._SCHEMA
 
-    def _write_raw_data(self, name, data):
+    def _write_raw_data(self, name: str, data: bytes | bytearray) -> None:
         '''
         Helper function to write data
 
@@ -98,7 +98,7 @@ def _write_raw_data(self, name, data):
         storage.write(data)
         storage.close()
 
-    def _write_schema(self, name):
+    def _write_schema(self, name: str) -> None:
         '''
         Writes resource schema
 
@@ -107,20 +107,20 @@ def _write_schema(self, name):
         self._write_raw_data(f"{name}.schema", bytes(
             self._RESOURCES[name].schema, 'utf-8'))
 
-    def _write_archive_signature(self):
+    def _write_archive_signature(self) -> None:
         '''Writes archive's signature'''
         self._write_raw_data(f"{self._NAME}.archive", b'\x00' * 16)
 
-    def _write_archive_schema(self):
+    def _write_archive_schema(self) -> None:
         '''Writes archive schema'''
         self._write_raw_data(
             f"{self._NAME}.archive.schema", bytes(self._SCHEMA, 'utf-8'))
 
-    def _write_index_schema(self, resource_name, schema):
+    def _write_index_schema(self, resource_name: str, schema: str) -> None:
         self._write_raw_data(
             f"{resource_name}_index.schema", bytes(schema, 'utf-8'))
 
-    def subarchive(self, name):
+    def subarchive(self, name: str) -> None:
         """
         Returns an archive builder for the sub-archive `name`.
         :raises $name_not_subarchive_error
@@ -129,7 +129,7 @@ def subarchive(self, name):
         NotImplemented
 
     @classmethod
-    def __validate_structure_fields(cls, name, struct, initializer):
+    def __validate_structure_fields(cls, name: str, struct: dict[str, Any], initializer: Any) -> None:
         '''
         Validates whether passed object has all required fields
 
@@ -146,7 +146,7 @@ def __validate_structure_fields(cls, name, struct, initializer):
             if key not in initializer._FIELD_KEYS:
                 raise UnknownFieldError(key, name)
 
-    def __set_instance(self, storage, name, value):
+    def __set_instance(self, storage: Any, name: str, value: dict[str, Any]) -> None:
         '''
         Creates and writes instance type resource
 
@@ -164,7 +164,7 @@ def __set_instance(self, storage, name, value):
 
         storage.write(bout)
 
-    def __set_vector(self, storage, name, vector):
+    def __set_vector(self, storage: Any, name: str, vector: list[dict[str, Any]]) -> None:
         '''
         Creates and writes vector resource
 
@@ -183,7 +183,7 @@ def __set_vector(self, storage, name, vector):
                             field.is_signed, value[key])
             storage.write(bout)
 
-    def __set_multivector(self, storage, name, value):
+    def __set_multivector(self, storage: Any, name: str, value: list[list[dict[str, Any]]]) -> None:
         '''
         Creates and writes multivector resource
 
@@ -252,7 +252,7 @@ def validate_fields(_obj):
         self._resources_written.append(name)
         self._resources_written.append(f'{name}_index')
 
-    def set(self, name, value):
+    def set(self, name: str, value: Any) -> None:
         """
         Write a resource for this archive at once.
         Can only be done once. `set` and `start` can't be used for the same resource.
@@ -288,7 +288,7 @@ def set(self, name, value):
 
         self._resources_written.append(name)
 
-    def finish(self):
+    def finish(self) -> None:
         """
         Closes the storage manager
         """

flatdata-py/flatdata/lib/data_access.py

@@ -3,13 +3,17 @@
  See the LICENSE file in the root of this project for license details.
 '''
 
+from collections.abc import Callable
+from typing import Any
+
 import numpy as np
+from numpy.typing import NDArray
 
 # Sign bits cache for the value reading.
 _SIGN_BITS = [0] + [(1 << (bits - 1)) for bits in range(1, 65)]
 
 
-def make_field_reader(offset_bits, num_bits, is_signed):
+def make_field_reader(offset_bits: int, num_bits: int, is_signed: bool) -> Callable[[Any, int], int]:
     """Build a specialized closure for reading a single field from a structure.
 
     Returns a function reader(data, pos_bytes) that reads the field value
@@ -76,7 +80,7 @@ def reader(data, pos):
     return reader
 
 
-def read_field_vectorized(raw_bytes_2d, field_offset_bits, field_width_bits, is_signed):
+def read_field_vectorized(raw_bytes_2d: NDArray[np.uint8], field_offset_bits: int, field_width_bits: int, is_signed: bool) -> NDArray[Any]:
     """Read a bit-packed field from all elements at once, returning a numpy array.
 
     :param raw_bytes_2d: numpy uint8 array shaped (num_elements, struct_size_bytes)
@@ -122,7 +126,7 @@ def read_field_vectorized(raw_bytes_2d, field_offset_bits, field_width_bits, is_
     return result
 
 
-def read_value(data, offset_bits, num_bits, is_signed):
+def read_value(data: Any, offset_bits: int, num_bits: int, is_signed: bool) -> int:
     """Read a bit-packed value from data at the given bit offset.
 
     This is a convenience wrapper around :func:`make_field_reader` for one-off
@@ -133,7 +137,7 @@ def read_value(data, offset_bits, num_bits, is_signed):
     return reader(data, 0)
 
 
-def write_value(data, offset_bits, num_bits, is_signed, value):
+def write_value(data: bytearray, offset_bits: int, num_bits: int, is_signed: bool, value: int) -> None:
     assert num_bits <= 64, f'Number of bits to write is greater than 64'
 
     offset_bytes, offset_extra_bits = divmod(offset_bits, 8)