Typed adapters (bluesky#1047)

* fix(devcontainer): Add stage to Docker build for devcontainer

* chore(types): Add type hints and hierarchy to Structure types

* fix type checking issues on various versions

* ClassVar

* Numpy types

* Add root Adapter type

* chore: Type hint and rationalise adapters

* Changes to fix pre-commit

* Fixes for pre-commit including correct typing of MapAdapter

* Fixes for pre-commit

* Adopt Zarr expectations of JSON type

* Extract Zarr Adapter from array assumptions

* Fixes for pre-commit

* Retype xarray adapters, MadAdapter

* Add types-cachetools to dev requirements

* Fixes to pass pre-commit

* Remove type vars from non generic classes

* Add SQLAdapter.structure_family

* Fix treating .structure_family as a property

* Make ParquetDatabaseAdapter inherit from Adapter[TableStructure]

* Fixes for pre-commit

* Make structure_family an attribute, make metadata a callable

* Make CSVAdapter inherit directly from Adapter[TableStructure]

* Fix structure family pickling issue

* Fix hdf5 adapter self referential metadata

* Change ZarrGroupAdapter conatiner structure to use list of keys

* Change NPYAdapter to inherit directly from Adapter[ArrayStructure]

* Remove unused print statement

* Remove AnyStructure

* Fix typing for backwards compatibility

* Fix metadata argument

* Add type to ExceAdapter(MadAdapter) generic

* Add return type, remove unused type ignore

* Fix JPEGAdapter.read_block signature, remove unused type ignore

* Fix merge conflict resolution

---------

Co-authored-by: Daniel Fernandes <[email protected]>
Co-authored-by: Dan Allan <[email protected]>

Author: 3 people

.pre-commit-config.yaml

@@ -35,3 +35,4 @@ repos:
     hooks:
     -    id: mypy
          args: [--strict]
+         additional_dependencies: [types-cachetools]

pyproject.toml

@@ -169,6 +169,7 @@ dev = [
     "sphinx-copybutton",
     "sphinx_rtd_theme",
     "sphinxcontrib-mermaid",
+    "types-cachetools",
 ]
 # These are used by the server to read files in the respective formats.
 formats = [

tiled/adapters/array.py

@@ -1,19 +1,20 @@
 import contextlib
-from typing import Any, List, Optional, Set, Tuple
+from typing import Any, List, Optional, Tuple
 
 import dask.array
 import numpy
 import pandas
 from numpy.typing import NDArray
 
+from tiled.adapters.core import Adapter
+
 from ..ndslice import NDSlice
-from ..storage import Storage
 from ..structures.array import ArrayStructure
 from ..structures.core import Spec, StructureFamily
 from ..type_aliases import JSON
 
 
-class ArrayAdapter:
+class ArrayAdapter(Adapter[ArrayStructure]):
     """
     Wrap an array-like object in an interface that Tiled can serve.
 
@@ -29,7 +30,6 @@ class ArrayAdapter:
     """
 
     structure_family = StructureFamily.array
-    supported_storage: Set[type[Storage]] = set()
 
     def __init__(
         self,
@@ -39,19 +39,8 @@ def __init__(
         metadata: Optional[JSON] = None,
         specs: Optional[List[Spec]] = None,
     ) -> None:
-        """
-
-        Parameters
-        ----------
-        array :
-        structure :
-        metadata :
-        specs :
-        """
         self._array = array
-        self._structure = structure
-        self._metadata = metadata or {}
-        self.specs = specs or []
+        super().__init__(structure, metadata=metadata, specs=specs)
 
     @classmethod
     def from_array(
@@ -64,21 +53,6 @@ def from_array(
         metadata: Optional[JSON] = None,
         specs: Optional[List[Spec]] = None,
     ) -> "ArrayAdapter":
-        """
-
-        Parameters
-        ----------
-        array :
-        shape :
-        chunks :
-        dims :
-        metadata :
-        specs :
-
-        Returns
-        -------
-
-        """
         # May be a list of something; convert to array
         if not hasattr(array, "__array__"):
             array = numpy.asanyarray(array)
@@ -118,26 +92,10 @@ def __repr__(self) -> str:
     def dims(self) -> Optional[Tuple[str, ...]]:
         return self._structure.dims
 
-    def metadata(self) -> JSON:
-        return self._metadata
-
-    def structure(self) -> ArrayStructure:
-        return self._structure
-
     def read(
         self,
         slice: NDSlice = NDSlice(...),
     ) -> NDArray[Any]:
-        """
-
-        Parameters
-        ----------
-        slice :
-
-        Returns
-        -------
-
-        """
         # _array[...] requires an actual tuple, not just a subclass of tuple
         array = self._array[tuple(slice)] if slice else self._array
         if isinstance(self._array, dask.array.Array):
@@ -149,17 +107,6 @@ def read_block(
         block: Tuple[int, ...],
         slice: NDSlice = NDSlice(...),
     ) -> NDArray[Any]:
-        """
-
-        Parameters
-        ----------
-        block :
-        slice :
-
-        Returns
-        -------
-
-        """
         # Slice the whole array to get this block.
         slice_, _ = slice_and_shape_from_block_and_chunks(block, self._structure.chunks)
         # _array[...] requires an actual tuple, not just a subclass of tuple
@@ -176,14 +123,6 @@ def slice_and_shape_from_block_and_chunks(
 ) -> tuple[NDSlice, NDSlice]:
     """
     Given dask-like chunks and block id, return slice and shape of the block.
-    Parameters
-    ----------
-    block :
-    chunks :
-
-    Returns
-    -------
-
     """
     slice_ = []
     shape = []

tiled/adapters/arrow.py

@@ -1,12 +1,14 @@
 import copy
+from collections.abc import Set
 from pathlib import Path
 from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple, Union
 from urllib.parse import quote_plus
 
 import pandas
 import pyarrow
 import pyarrow.feather as feather
-import pyarrow.fs
+
+from tiled.adapters.core import Adapter
 
 from ..catalog.orm import Node
 from ..storage import FileStorage, Storage
@@ -19,37 +21,29 @@
 from .utils import init_adapter_from_catalog
 
 
-class ArrowAdapter:
+class ArrowAdapter(Adapter[TableStructure]):
     """ArrowAdapter Class"""
 
-    structure_family = StructureFamily.table
-    supported_storage = {FileStorage}
+    structure_family: StructureFamily = StructureFamily.table
 
     def __init__(
         self,
         data_uris: List[str],
         structure: Optional[TableStructure] = None,
+        *,
         metadata: Optional[JSON] = None,
         specs: Optional[List[Spec]] = None,
-        **kwargs: Optional[Any],
     ) -> None:
-        """
-
-        Parameters
-        ----------
-        data_uris : list of uris where data sits.
-        structure :
-        metadata :
-        specs :
-        """
         # TODO Store data_uris instead and generalize to non-file schemes.
         self._partition_paths = [path_from_uri(uri) for uri in data_uris]
-        self._metadata = metadata or {}
         if structure is None:
             table = feather.read_table(self._partition_paths)
             structure = TableStructure.from_arrow_table(table)
-        self._structure = structure
-        self.specs = list(specs or [])
+        super().__init__(structure, metadata=metadata, specs=specs)
+
+    @classmethod
+    def supported_storage(cls) -> Set[type[Storage]]:
+        return {FileStorage}
 
     @classmethod
     def from_catalog(
@@ -59,10 +53,7 @@ def from_catalog(
         /,
         **kwargs: Optional[Any],
     ) -> "ArrowAdapter":
-        return init_adapter_from_catalog(cls, data_source, node, **kwargs)  # type: ignore
-
-    def metadata(self) -> JSON:
-        return self._metadata
+        return init_adapter_from_catalog(cls, data_source, node, **kwargs)
 
     @classmethod
     def init_storage(
@@ -101,9 +92,6 @@ def init_storage(
         data_source.assets.extend(assets)
         return data_source
 
-    def structure(self) -> TableStructure:
-        return self._structure
-
     def get(self, key: str) -> Union[ArrayAdapter, None]:
         if key not in self.structure().columns:
             return None
@@ -153,6 +141,7 @@ def from_single_file(
         cls,
         data_uri: str,
         structure: Optional[TableStructure] = None,
+        *,
         metadata: Optional[JSON] = None,
         specs: Optional[List[Spec]] = None,
     ) -> "ArrowAdapter":

tiled/adapters/awkward.py

@@ -4,21 +4,22 @@
 import awkward.forms
 from numpy.typing import NDArray
 
-from ..storage import FileStorage
+from tiled.adapters.core import Adapter
+
 from ..structures.awkward import AwkwardStructure
 from ..structures.core import Spec, StructureFamily
 from ..type_aliases import JSON
 from .awkward_directory_container import DirectoryContainer
 
 
-class AwkwardAdapter:
-    structure_family = StructureFamily.awkward
-    supported_storage = {FileStorage}
+class AwkwardAdapter(Adapter[AwkwardStructure]):
+    structure_family: StructureFamily = StructureFamily.awkward
 
     def __init__(
         self,
         container: DirectoryContainer,
         structure: AwkwardStructure,
+        *,
         metadata: Optional[JSON] = None,
         specs: Optional[List[Spec]] = None,
     ) -> None:
@@ -32,14 +33,13 @@ def __init__(
         specs :
         """
         self.container = container
-        self._metadata = metadata or {}
-        self._structure = structure
-        self.specs = list(specs or [])
+        super().__init__(structure, metadata=metadata, specs=specs)
 
     @classmethod
     def from_array(
         cls,
         array: NDArray[Any],
+        *,
         metadata: Optional[JSON] = None,
         specs: Optional[List[Spec]] = None,
     ) -> "AwkwardAdapter":
@@ -64,15 +64,6 @@ def from_array(
             specs=specs,
         )
 
-    def metadata(self) -> JSON:
-        """
-
-        Returns
-        -------
-
-        """
-        return self._metadata
-
     def read_buffers(self, form_keys: Optional[List[str]] = None) -> Dict[str, bytes]:
         """
 
@@ -118,6 +109,3 @@ def write(self, container: DirectoryContainer) -> None:
         """
         for form_key, value in container.items():
             self.container[form_key] = value
-
-    def structure(self) -> AwkwardStructure:
-        return self._structure

tiled/adapters/awkward_buffers.py

@@ -3,6 +3,7 @@
 """
 
 import copy
+from collections.abc import Set
 from pathlib import Path
 from typing import Any, List, Optional
 from urllib.parse import quote_plus
@@ -12,7 +13,7 @@
 from ..catalog.orm import Node
 from ..storage import FileStorage, Storage
 from ..structures.awkward import AwkwardStructure
-from ..structures.core import Spec, StructureFamily
+from ..structures.core import Spec
 from ..structures.data_source import Asset, DataSource
 from ..type_aliases import JSON
 from ..utils import path_from_uri
@@ -22,9 +23,6 @@
 
 
 class AwkwardBuffersAdapter(AwkwardAdapter):
-    structure_family = StructureFamily.awkward
-    supported_storage = {FileStorage}
-
     @classmethod
     def init_storage(
         cls,
@@ -58,6 +56,7 @@ def __init__(
         self,
         data_uri: str,
         structure: AwkwardStructure,
+        *,
         metadata: Optional[JSON] = None,
         specs: Optional[List[Spec]] = None,
     ):
@@ -73,6 +72,10 @@ def __init__(
             specs=specs,
         )
 
+    @classmethod
+    def supported_storage(cls) -> Set[type[Storage]]:
+        return {FileStorage}
+
     @classmethod
     def from_catalog(
         cls,
@@ -81,4 +84,4 @@ def from_catalog(
         /,
         **kwargs: Optional[Any],
     ) -> "AwkwardBuffersAdapter":
-        return init_adapter_from_catalog(cls, data_source, node, **kwargs)  # type: ignore
+        return init_adapter_from_catalog(cls, data_source, node, **kwargs)

tiled/adapters/container.py

@@ -0,0 +1,10 @@
+from collections.abc import Mapping
+from typing import Generic
+
+from tiled.adapters.core import A, Adapter
+from tiled.structures.container import ContainerStructure
+from tiled.structures.core import StructureFamily
+
+
+class ContainerAdapter(Adapter[ContainerStructure], Mapping[str, A], Generic[A]):
+    structure_family: StructureFamily = StructureFamily.container