chore: Add types to Structures

Author: DiamondJoseph

CHANGELOG.md

@@ -8,6 +8,7 @@ Write the date in place of the "Unreleased" in the case a new version is release
 ### Changed
 
 - Typehint utils collection implementations
+- Typehint Structure types
 
 ## v0.1.0-b35 (2025-08-20)
 

tiled/client/base.py

@@ -2,7 +2,7 @@
 from copy import copy, deepcopy
 from dataclasses import asdict
 from pathlib import Path
-from typing import Dict, List, Union
+from typing import Dict, List, Optional, Union
 from urllib.parse import parse_qs, urlparse
 
 import json_merge_patch
@@ -11,6 +11,7 @@
 from httpx import URL
 
 from tiled.client.context import Context
+from tiled.structures.root import Structure
 
 from ..structures.core import STRUCTURE_TYPES, Spec, StructureFamily
 from ..structures.data_source import DataSource
@@ -131,8 +132,8 @@ def __init__(
         *,
         item,
         structure_clients,
-        structure=None,
-        include_data_sources=False,
+        structure: Optional[Structure] = None,
+        include_data_sources: bool = False,
     ):
         self._context = context
         self._item = item

tiled/server/schemas.py

@@ -16,6 +16,8 @@
 from pydantic_core import PydanticCustomError
 from typing_extensions import Annotated, TypedDict
 
+from tiled.structures.root import Structure
+
 from ..structures.array import ArrayStructure
 from ..structures.awkward import AwkwardStructure
 from ..structures.core import STRUCTURE_TYPES, StructureFamily
@@ -30,7 +32,7 @@
 DataT = TypeVar("DataT")
 LinksT = TypeVar("LinksT")
 MetaT = TypeVar("MetaT")
-StructureT = TypeVar("StructureT")
+StructureT = TypeVar("StructureT", bound=Structure)
 
 
 MAX_ALLOWED_SPECS = 20

tiled/structures/array.py

@@ -1,11 +1,18 @@
 import enum
 import os
 import sys
+from collections.abc import Mapping
 from dataclasses import dataclass
-from typing import List, Optional, Tuple, Union
+from typing import Any, ClassVar, List, Optional, Tuple, Union
 
 import numpy
 
+from tiled.structures.root import Structure
+
+# from dtype.descr
+FieldDescr = Union[Tuple[str, str], Tuple[str, str, Tuple[int, ...]]]
+NumpyDescr = List[FieldDescr]
+
 
 class Endianness(str, enum.Enum):
     """
@@ -57,7 +64,7 @@ class Kind(str, enum.Enum):
         object = "O"  # Object (i.e. the memory contains a pointer to PyObject)
 
     @classmethod
-    def _missing_(cls, key):
+    def _missing_(cls, key: str):
         if key == "O":
             raise ObjectArrayTypeDisabled(
                 "Numpy 'object'-type arrays are not enabled by default "
@@ -78,17 +85,19 @@ class BuiltinDtype:
     itemsize: int
     dt_units: Optional[str] = None
 
-    __endianness_map = {
+    __endianness_map: ClassVar[Mapping[str, str]] = {
         ">": "big",
         "<": "little",
         "=": sys.byteorder,
         "|": "not_applicable",
     }
 
-    __endianness_reverse_map = {"big": ">", "little": "<", "not_applicable": "|"}
+    __endianness_reverse_map: ClassVar[Mapping[str, str]] = {
+        v: k for k, v in __endianness_map.items() if k != "="
+    }
 
     @classmethod
-    def from_numpy_dtype(cls, dtype) -> "BuiltinDtype":
+    def from_numpy_dtype(cls, dtype: numpy.dtype) -> "BuiltinDtype":
         # Extract datetime units from the dtype string representation,
         # e.g. `'<M8[ns]'` has `dt_units = '[ns]'`. Count determines the number of base units in a step.
         dt_units = None
@@ -106,7 +115,7 @@ def from_numpy_dtype(cls, dtype) -> "BuiltinDtype":
     def to_numpy_dtype(self) -> numpy.dtype:
         return numpy.dtype(self.to_numpy_str())
 
-    def to_numpy_str(self):
+    def to_numpy_str(self) -> str:
         endianness = self.__endianness_reverse_map[self.endianness]
         # dtype.itemsize always reports bytes.  The format string from the
         # numeric types the string format is: {type_code}{byte_count} so we can
@@ -125,7 +134,7 @@ def to_numpy_descr(self):
         return self.to_numpy_str()
 
     @classmethod
-    def from_json(cls, structure):
+    def from_json(cls, structure: Mapping[str, Any]) -> "BuiltinDtype":
         return cls(
             kind=Kind(structure["kind"]),
             itemsize=structure["itemsize"],
@@ -141,7 +150,7 @@ class Field:
     shape: Optional[Tuple[int, ...]]
 
     @classmethod
-    def from_numpy_descr(cls, field):
+    def from_numpy_descr(cls, field: FieldDescr) -> "Field":
         name, *rest = field
         if name == "":
             raise ValueError(
@@ -159,7 +168,7 @@ def from_numpy_descr(cls, field):
             FType = StructDtype.from_numpy_dtype(numpy.dtype(f_type))
         return cls(name=name, dtype=FType, shape=shape)
 
-    def to_numpy_descr(self):
+    def to_numpy_descr(self) -> FieldDescr:
         if isinstance(self.dtype, BuiltinDtype):
             base = [self.name, self.dtype.to_numpy_str()]
         else:
@@ -170,7 +179,7 @@ def to_numpy_descr(self):
             return tuple(base + [self.shape])
 
     @classmethod
-    def from_json(cls, structure):
+    def from_json(cls, structure: Mapping[str, Any]) -> "Field":
         name = structure["name"]
         if "fields" in structure["dtype"]:
             ftype = StructDtype.from_json(structure["dtype"])
@@ -185,7 +194,7 @@ class StructDtype:
     fields: List[Field]
 
     @classmethod
-    def from_numpy_dtype(cls, dtype):
+    def from_numpy_dtype(cls, dtype: numpy.dtype) -> "StructDtype":
         # subdtypes push extra dimensions into arrays, we should handle these
         # a layer up and report an array with bigger dimensions.
         if dtype.subdtype is not None:
@@ -198,36 +207,36 @@ def from_numpy_dtype(cls, dtype):
             fields=[Field.from_numpy_descr(f) for f in dtype.descr],
         )
 
-    def to_numpy_dtype(self):
+    def to_numpy_dtype(self) -> numpy.dtype:
         return numpy.dtype(self.to_numpy_descr())
 
-    def to_numpy_descr(self):
+    def to_numpy_descr(self) -> NumpyDescr:
         return [f.to_numpy_descr() for f in self.fields]
 
-    def max_depth(self):
+    def max_depth(self) -> int:
         return max(
             1 if isinstance(f.dtype, BuiltinDtype) else 1 + f.dtype.max_depth()
             for f in self.fields
         )
 
     @classmethod
-    def from_json(cls, structure):
+    def from_json(cls, structure: Mapping[str, Any]) -> "StructDtype":
         return cls(
             itemsize=structure["itemsize"],
             fields=[Field.from_json(f) for f in structure["fields"]],
         )
 
 
 @dataclass
-class ArrayStructure:
+class ArrayStructure(Structure):
     data_type: Union[BuiltinDtype, StructDtype]
     chunks: Tuple[Tuple[int, ...], ...]  # tuple-of-tuples-of-ints like ((3,), (3,))
     shape: Tuple[int, ...]  # tuple of ints like (3, 3)
     dims: Optional[Tuple[str, ...]] = None  # None or tuple of names like ("x", "y")
     resizable: Union[bool, Tuple[bool, ...]] = False
 
     @classmethod
-    def from_json(cls, structure):
+    def from_json(cls, structure: Mapping[str, Any]) -> "ArrayStructure":
         if "fields" in structure["data_type"]:
             data_type = StructDtype.from_json(structure["data_type"])
         else:
@@ -273,7 +282,7 @@ def from_array(cls, array, shape=None, chunks=None, dims=None) -> "ArrayStructur
             data_type = StructDtype.from_numpy_dtype(array.dtype)
         else:
             data_type = BuiltinDtype.from_numpy_dtype(array.dtype)
-        return ArrayStructure(
+        return cls(
             data_type=data_type,
             shape=shape,
             chunks=normalized_chunks,

tiled/structures/awkward.py

@@ -1,19 +1,24 @@
 from dataclasses import dataclass
+from typing import Any, Iterable, Mapping, Optional
 
 import awkward
 
+from tiled.structures.root import Structure
+
 
 @dataclass
-class AwkwardStructure:
+class AwkwardStructure(Structure):
     length: int
     form: dict
 
     @classmethod
-    def from_json(cls, structure):
+    def from_json(cls, structure: Mapping[str, Any]) -> "AwkwardStructure":
         return cls(**structure)
 
 
-def project_form(form, form_keys_touched):
+def project_form(
+    form: awkward.forms.Form, form_keys_touched: Iterable[str]
+) -> Optional[awkward.forms.Form]:
     # See https://github.com/bluesky/tiled/issues/450
     if isinstance(form, awkward.forms.RecordForm):
         if form.fields is None:

tiled/structures/core.py

@@ -9,6 +9,8 @@
 from dataclasses import asdict, dataclass
 from typing import Dict, Optional
 
+from tiled.structures.root import Structure
+
 from ..utils import OneShotCachedMap
 
 
@@ -47,8 +49,7 @@ def dict(self) -> Dict[str, Optional[str]]:
     model_dump = dict  # For easy interoperability with pydantic 2.x models
 
 
-# TODO: make type[Structure] after #1036
-STRUCTURE_TYPES = OneShotCachedMap[StructureFamily, type](
+STRUCTURE_TYPES = OneShotCachedMap[StructureFamily, type[Structure]](
     {
         StructureFamily.array: lambda: importlib.import_module(
             "...structures.array", StructureFamily.__module__

tiled/structures/data_source.py

@@ -1,6 +1,9 @@
 import dataclasses
 import enum
-from typing import Generic, List, Optional, TypeVar
+from collections.abc import Mapping
+from typing import Any, Generic, List, Optional, TypeVar
+
+from tiled.structures.root import Structure
 
 from .core import StructureFamily
 
@@ -21,21 +24,21 @@ class Asset:
     id: Optional[int] = None
 
 
-StructureT = TypeVar("StructureT")
+StructureT = TypeVar("StructureT", bound=Structure)
 
 
 @dataclasses.dataclass
 class DataSource(Generic[StructureT]):
     structure_family: StructureFamily
-    structure: StructureT
+    structure: Optional[StructureT]
     id: Optional[int] = None
     mimetype: Optional[str] = None
     parameters: dict = dataclasses.field(default_factory=dict)
     assets: List[Asset] = dataclasses.field(default_factory=list)
     management: Management = Management.writable
 
     @classmethod
-    def from_json(cls, d):
-        d = d.copy()
+    def from_json(cls, structure: Mapping[str, Any]) -> "DataSource":
+        d = structure.copy()
         assets = [Asset(**a) for a in d.pop("assets")]
         return cls(assets=assets, **d)

tiled/structures/root.py

@@ -0,0 +1,12 @@
+import dataclasses
+from abc import ABC
+from collections.abc import Mapping
+from typing import Any
+
+
+@dataclasses.dataclass
+class Structure(ABC):
+    @classmethod
+    # TODO: When dropping support for Python 3.10 replace with -> Self
+    def from_json(cls, structure: Mapping[str, Any]) -> "Structure":
+        ...

tiled/structures/sparse.py

@@ -1,6 +1,9 @@
 import enum
+from collections.abc import Mapping
 from dataclasses import dataclass, field
-from typing import Optional, Tuple, Union
+from typing import Any, Optional, Tuple, Union
+
+from tiled.structures.root import Structure
 
 from .array import BuiltinDtype, Endianness, Kind, StructDtype
 
@@ -12,7 +15,7 @@ class SparseLayout(str, enum.Enum):
 
 
 @dataclass
-class COOStructure:
+class COOStructure(Structure):
     chunks: Tuple[Tuple[int, ...], ...]  # tuple-of-tuples-of-ints like ((3,), (3,))
     shape: Tuple[int, ...]  # tuple of ints like (3, 3)
     data_type: Optional[Union[BuiltinDtype, StructDtype]] = None
@@ -27,7 +30,7 @@ class COOStructure:
     # TODO Include fill_value?
 
     @classmethod
-    def from_json(cls, structure):
+    def from_json(cls, structure: Mapping[str, Any]) -> "COOStructure":
         data_type = structure.get("data_type", None)
         if data_type is not None and "fields" in data_type:
             data_type = StructDtype.from_json(data_type)