#152 Update generic type hints

Author: astropenguin

pyproject.toml

@@ -36,12 +36,10 @@ pytest = "^7.1"
 sphinx = "^4.5"
 
 [tool.pyright]
-typeCheckingMode = "strict"
 reportImportCycles = "warning"
 reportUnknownArgumentType = "warning"
 reportUnknownMemberType = "warning"
-reportUnknownParameterType = "warning"
-reportUnknownVariableType = "warning"
+typeCheckingMode = "strict"
 
 [build-system]
 requires = ["poetry-core>=1.0.0"]

xarray_dataclasses/__init__.py

@@ -2,18 +2,6 @@
 __version__ = "1.1.0"
 
 
-# for Python 3.7 and 3.8
-def _make_field_generic():
-    from dataclasses import Field
-    from typing import Sequence
-
-    GenericAlias = type(Sequence[int])
-    Field.__class_getitem__ = classmethod(GenericAlias)
-
-
-_make_field_generic()
-
-
 # submodules
 from . import dataarray
 from . import dataset

xarray_dataclasses/dataarray.py

@@ -27,7 +27,7 @@
 # submodules
 from .datamodel import DataModel
 from .dataoptions import DataOptions
-from .typing import DataClass, DataClassFields, DataType, Order, Shape, Sizes
+from .typing import AnyArray, DataClass, DataClassFields, DataType, Order, Shape, Sizes
 
 
 # type hints
@@ -168,7 +168,7 @@ def new(cls: Any, *args: Any, **kwargs: Any) -> Any:
     @classmethod
     def shaped(
         cls: Type[OptionedClass[PInit, TDataArray]],
-        func: Callable[[Shape], np.ndarray],
+        func: Callable[[Shape], AnyArray],
         shape: Union[Shape, Sizes],
         **kwargs: Any,
     ) -> TDataArray:
@@ -178,7 +178,7 @@ def shaped(
     @classmethod
     def shaped(
         cls: Type[DataClass[PInit]],
-        func: Callable[[Shape], np.ndarray],
+        func: Callable[[Shape], AnyArray],
         shape: Union[Shape, Sizes],
         **kwargs: Any,
     ) -> xr.DataArray:
@@ -187,7 +187,7 @@ def shaped(
     @classmethod
     def shaped(
         cls: Any,
-        func: Callable[[Shape], np.ndarray],
+        func: Callable[[Shape], AnyArray],
         shape: Union[Shape, Sizes],
         **kwargs: Any,
     ) -> Any:

xarray_dataclasses/datamodel.py

@@ -3,7 +3,7 @@
 
 
 # standard library
-from dataclasses import Field, dataclass, field, is_dataclass
+from dataclasses import dataclass, field, is_dataclass
 from typing import Any, Dict, Hashable, List, Optional, Tuple, Type, Union, cast
 
 
@@ -15,6 +15,7 @@
 
 # submodules
 from .typing import (
+    AnyField,
     DataClass,
     DataType,
     Dims,
@@ -204,7 +205,7 @@ def eval_dataclass(dataclass: AnyDataClass[PInit]) -> None:
         field.type = types[field.name]
 
 
-def get_entry(field: Field[Any], value: Any) -> AnyEntry:
+def get_entry(field: AnyField, value: Any) -> AnyEntry:
     """Create an entry from a field and its value."""
     field_type = get_field_type(field.type)
     repr_type = get_repr_type(field.type)

xarray_dataclasses/dataset.py

@@ -18,7 +18,7 @@
 # submodules
 from .datamodel import DataModel
 from .dataoptions import DataOptions
-from .typing import DataClass, DataClassFields, DataType, Order, Shape, Sizes
+from .typing import AnyArray, DataClass, DataClassFields, DataType, Order, Shape, Sizes
 
 
 # type hints
@@ -159,7 +159,7 @@ def new(cls: Any, *args: Any, **kwargs: Any) -> Any:
     @classmethod
     def shaped(
         cls: Type[OptionedClass[PInit, TDataset]],
-        func: Callable[[Shape], np.ndarray],
+        func: Callable[[Shape], AnyArray],
         sizes: Sizes,
         **kwargs: Any,
     ) -> TDataset:
@@ -169,7 +169,7 @@ def shaped(
     @classmethod
     def shaped(
         cls: Type[DataClass[PInit]],
-        func: Callable[[Shape], np.ndarray],
+        func: Callable[[Shape], AnyArray],
         sizes: Sizes,
         **kwargs: Any,
     ) -> xr.Dataset:
@@ -178,7 +178,7 @@ def shaped(
     @classmethod
     def shaped(
         cls: Any,
-        func: Callable[[Shape], np.ndarray],
+        func: Callable[[Shape], AnyArray],
         sizes: Sizes,
         **kwargs: Any,
     ) -> Any:
@@ -194,7 +194,7 @@ def shaped(
 
         """
         model = DataModel.from_dataclass(cls)
-        data_vars: Dict[str, np.ndarray] = {}
+        data_vars: Dict[str, AnyArray] = {}
 
         for key, entry in model.data_vars_items:
             shape = tuple(sizes[dim] for dim in entry.dims)

xarray_dataclasses/typing.py

@@ -35,13 +35,15 @@
 
 
 # dependencies
+import numpy as np
 import xarray as xr
 from more_itertools import collapse
 from typing_extensions import (
     Annotated,
     Literal,
     ParamSpec,
     Protocol,
+    TypeAlias,
     get_args,
     get_origin,
     get_type_hints,
@@ -56,7 +58,9 @@
 TDtype = TypeVar("TDtype", covariant=True)
 TName = TypeVar("TName", bound=Hashable)
 
-DataClassFields = Dict[str, Field[Any]]
+AnyArray: TypeAlias = "np.ndarray[Any, Any]"
+AnyField: TypeAlias = "Field[Any]"
+DataClassFields = Dict[str, AnyField]
 DataType = Union[xr.DataArray, xr.Dataset]
 Dims = Tuple[str, ...]
 Dtype = Optional[str]