#135 Make typing.DataClass generic

Author: astropenguin

xarray_dataclasses/dataarray.py

@@ -2,10 +2,18 @@
 
 
 # standard library
-from dataclasses import Field
 from functools import partial, wraps
 from types import MethodType
-from typing import Any, Callable, Dict, Optional, Type, TypeVar, Union, overload
+from typing import (
+    Any,
+    Callable,
+    ClassVar,
+    Optional,
+    Type,
+    TypeVar,
+    Union,
+    overload,
+)
 
 
 # dependencies
@@ -18,7 +26,7 @@
 # submodules
 from .datamodel import DataModel
 from .dataoptions import DataOptions
-from .typing import DataType, Order, Shape, Sizes
+from .typing import DataClass, DataClassFields, DataType, Order, Shape, Sizes
 
 
 # constants
@@ -30,29 +38,20 @@
 TDataArray = TypeVar("TDataArray", bound=xr.DataArray)
 
 
-class DataClass(Protocol[P]):
-    """Type hint for a dataclass object."""
+class OptionedClass(Protocol[P, TDataArray]):
+    """Type hint for dataclass objects with options."""
 
     def __init__(self, *args: P.args, **kwargs: P.kwargs) -> None:
         ...
 
-    __dataclass_fields__: Dict[str, Field[Any]]
-
-
-class DataArrayClass(Protocol[P, TDataArray]):
-    """Type hint for a dataclass object with a DataArray factory."""
-
-    def __init__(self, *args: P.args, **kwargs: P.kwargs) -> None:
-        ...
-
-    __dataclass_fields__: Dict[str, Field[Any]]
+    __dataclass_fields__: ClassVar[DataClassFields]
     __dataoptions__: DataOptions[TDataArray]
 
 
 # runtime functions
 @overload
 def asdataarray(
-    dataclass: DataArrayClass[Any, TDataArray],
+    dataclass: OptionedClass[P, TDataArray],
     reference: Optional[DataType] = None,
     dataoptions: DataOptions[Any] = DEFAULT_OPTIONS,
 ) -> TDataArray:
@@ -61,7 +60,7 @@ def asdataarray(
 
 @overload
 def asdataarray(
-    dataclass: DataClass[Any],
+    dataclass: DataClass[P],
     reference: Optional[DataType] = None,
     dataoptions: DataOptions[TDataArray] = DEFAULT_OPTIONS,
 ) -> TDataArray:
@@ -127,7 +126,7 @@ def __init__(self, func: Any) -> None:
     def __get__(
         self,
         obj: Any,
-        cls: Type[DataArrayClass[P, TDataArray]],
+        cls: Type[OptionedClass[P, TDataArray]],
     ) -> Callable[P, TDataArray]:
         ...
 
@@ -163,7 +162,7 @@ def new(cls: Any, *args: Any, **kwargs: Any) -> Any:
     @overload
     @classmethod
     def shaped(
-        cls: Type[DataArrayClass[P, TDataArray]],
+        cls: Type[OptionedClass[P, TDataArray]],
         func: Callable[[Shape], np.ndarray],
         shape: Union[Shape, Sizes],
         **kwargs: Any,
@@ -209,7 +208,7 @@ def shaped(
     @overload
     @classmethod
     def empty(
-        cls: Type[DataArrayClass[P, TDataArray]],
+        cls: Type[OptionedClass[P, TDataArray]],
         shape: Union[Shape, Sizes],
         order: Order = "C",
         **kwargs: Any,
@@ -251,7 +250,7 @@ def empty(
     @overload
     @classmethod
     def zeros(
-        cls: Type[DataArrayClass[P, TDataArray]],
+        cls: Type[OptionedClass[P, TDataArray]],
         shape: Union[Shape, Sizes],
         order: Order = "C",
         **kwargs: Any,
@@ -293,7 +292,7 @@ def zeros(
     @overload
     @classmethod
     def ones(
-        cls: Type[DataArrayClass[P, TDataArray]],
+        cls: Type[OptionedClass[P, TDataArray]],
         shape: Union[Shape, Sizes],
         order: Order = "C",
         **kwargs: Any,
@@ -335,7 +334,7 @@ def ones(
     @overload
     @classmethod
     def full(
-        cls: Type[DataArrayClass[P, TDataArray]],
+        cls: Type[OptionedClass[P, TDataArray]],
         shape: Union[Shape, Sizes],
         fill_value: Any,
         order: Order = "C",

xarray_dataclasses/datamodel.py

@@ -3,13 +3,13 @@
 
 # standard library
 from dataclasses import Field, dataclass, field, is_dataclass
-from typing import Any, Dict, Hashable, Optional, Type, cast
+from typing import Any, Dict, Hashable, Optional, Type, Union, cast
 
 
 # dependencies
 import numpy as np
 import xarray as xr
-from typing_extensions import TypedDict, get_type_hints
+from typing_extensions import ParamSpec, TypedDict, get_type_hints
 
 
 # submodules
@@ -28,6 +28,8 @@
 
 
 # type hints
+P = ParamSpec("P")
+AnyDataClass = Union[Type[DataClass[P]], DataClass[P]]
 DimsDtype = TypedDict("DimsDtype", dims=Dims, dtype=Dtype)
 
 
@@ -45,7 +47,7 @@ class Data:
     type: DimsDtype
     """Type (dims and dtype) of the field."""
 
-    factory: Optional[Type[DataClass]] = None
+    factory: Any = None
     """Factory dataclass to create a DataArray object."""
 
     def __call__(self, reference: Optional[DataType] = None) -> xr.DataArray:
@@ -137,7 +139,7 @@ class DataModel:
     """Model of the name fields."""
 
     @classmethod
-    def from_dataclass(cls, dataclass: DataClass) -> "DataModel":
+    def from_dataclass(cls, dataclass: AnyDataClass[P]) -> "DataModel":
         """Create a data model from a dataclass or its object."""
         model = cls()
         eval_field_types(dataclass)
@@ -162,7 +164,7 @@ def from_dataclass(cls, dataclass: DataClass) -> "DataModel":
 
 
 # runtime functions
-def eval_field_types(dataclass: DataClass) -> None:
+def eval_field_types(dataclass: AnyDataClass[P]) -> None:
     """Evaluate field types of a dataclass or its object."""
     hints = get_type_hints(dataclass, include_extras=True)  # type: ignore
 

xarray_dataclasses/dataset.py

@@ -2,10 +2,9 @@
 
 
 # standard library
-from dataclasses import Field
 from functools import partial, wraps
 from types import MethodType
-from typing import Any, Callable, Dict, Optional, Type, TypeVar, overload
+from typing import Any, Callable, ClassVar, Dict, Optional, Type, TypeVar, overload
 
 
 # dependencies
@@ -18,7 +17,7 @@
 # submodules
 from .datamodel import DataModel
 from .dataoptions import DataOptions
-from .typing import DataType, Order, Shape, Sizes
+from .typing import DataClass, DataClassFields, DataType, Order, Shape, Sizes
 
 
 # constants
@@ -30,29 +29,20 @@
 TDataset = TypeVar("TDataset", bound=xr.Dataset)
 
 
-class DataClass(Protocol[P]):
-    """Type hint for a dataclass object."""
+class OptionedClass(Protocol[P, TDataset]):
+    """Type hint for dataclass objects with options."""
 
     def __init__(self, *args: P.args, **kwargs: P.kwargs) -> None:
         ...
 
-    __dataclass_fields__: Dict[str, Field[Any]]
-
-
-class DatasetClass(Protocol[P, TDataset]):
-    """Type hint for a dataclass object with a Dataset factory."""
-
-    def __init__(self, *args: P.args, **kwargs: P.kwargs) -> None:
-        ...
-
-    __dataclass_fields__: Dict[str, Field[Any]]
+    __dataclass_fields__: ClassVar[DataClassFields]
     __dataoptions__: DataOptions[TDataset]
 
 
 # runtime functions and classes
 @overload
 def asdataset(
-    dataclass: DatasetClass[Any, TDataset],
+    dataclass: OptionedClass[P, TDataset],
     reference: Optional[DataType] = None,
     dataoptions: DataOptions[Any] = DEFAULT_OPTIONS,
 ) -> TDataset:
@@ -61,7 +51,7 @@ def asdataset(
 
 @overload
 def asdataset(
-    dataclass: DataClass[Any],
+    dataclass: DataClass[P],
     reference: Optional[DataType] = None,
     dataoptions: DataOptions[TDataset] = DEFAULT_OPTIONS,
 ) -> TDataset:
@@ -125,7 +115,7 @@ def __init__(self, func: Callable[..., Any]) -> None:
     def __get__(
         self,
         obj: Any,
-        cls: Type[DatasetClass[P, TDataset]],
+        cls: Type[OptionedClass[P, TDataset]],
     ) -> Callable[P, TDataset]:
         ...
 
@@ -161,7 +151,7 @@ def new(cls: Any, *args: Any, **kwargs: Any) -> Any:
     @overload
     @classmethod
     def shaped(
-        cls: Type[DatasetClass[P, TDataset]],
+        cls: Type[OptionedClass[P, TDataset]],
         func: Callable[[Shape], np.ndarray],
         sizes: Sizes,
         **kwargs: Any,
@@ -208,7 +198,7 @@ def shaped(
     @overload
     @classmethod
     def empty(
-        cls: Type[DatasetClass[P, TDataset]],
+        cls: Type[OptionedClass[P, TDataset]],
         sizes: Sizes,
         order: Order = "C",
         **kwargs: Any,
@@ -250,7 +240,7 @@ def empty(
     @overload
     @classmethod
     def zeros(
-        cls: Type[DatasetClass[P, TDataset]],
+        cls: Type[OptionedClass[P, TDataset]],
         sizes: Sizes,
         order: Order = "C",
         **kwargs: Any,
@@ -292,7 +282,7 @@ def zeros(
     @overload
     @classmethod
     def ones(
-        cls: Type[DatasetClass[P, TDataset]],
+        cls: Type[OptionedClass[P, TDataset]],
         sizes: Sizes,
         order: Order = "C",
         **kwargs: Any,
@@ -334,7 +324,7 @@ def ones(
     @overload
     @classmethod
     def full(
-        cls: Type[DatasetClass[P, TDataset]],
+        cls: Type[OptionedClass[P, TDataset]],
         sizes: Sizes,
         fill_value: Any,
         order: Order = "C",

xarray_dataclasses/typing.py

@@ -22,7 +22,7 @@
 from enum import auto, Enum
 from typing import (
     Any,
-    Callable,
+    ClassVar,
     Dict,
     Hashable,
     Optional,
@@ -38,6 +38,7 @@
 from typing_extensions import (
     Annotated,
     Literal,
+    ParamSpec,
     Protocol,
     get_args,
     get_origin,
@@ -74,21 +75,25 @@ def annotates(self, hint: Any) -> bool:
 
 
 # type hints
+P = ParamSpec("P")
+T = TypeVar("T")
+TDataClass = TypeVar("TDataClass", bound="DataClass[Any]")
+TDims = TypeVar("TDims", covariant=True)
+TDtype = TypeVar("TDtype", covariant=True)
+TName = TypeVar("TName", bound=Hashable)
+
+DataClassFields = Dict[str, Field[Any]]
 DataType = Union[xr.DataArray, xr.Dataset]
 Dims = Tuple[str, ...]
 Dtype = Optional[str]
 Order = Literal["C", "F"]
 Shape = Union[Sequence[int], int]
 Sizes = Dict[str, int]
-T = TypeVar("T")
-TDims = TypeVar("TDims", covariant=True)
-TDtype = TypeVar("TDtype", covariant=True)
-TName = TypeVar("TName", bound=Hashable)
 
 
 @runtime_checkable
 class ArrayLike(Protocol[TDims, TDtype]):
-    """Type hint of array-like objects."""
+    """Type hint for array-like objects."""
 
     def astype(self: T, dtype: Any) -> T:
         """Method to convert data type of the object."""
@@ -105,14 +110,13 @@ def shape(self) -> Tuple[int, ...]:
         ...
 
 
-class DataClass(Protocol):
-    """Type hint of dataclasses or their objects."""
-
-    __init__: Callable[..., None]
-    __dataclass_fields__: Dict[str, Field[Any]]
+class DataClass(Protocol[P]):
+    """Type hint for dataclass objects."""
 
+    def __init__(self, *args: P.args, **kwargs: P.kwargs) -> None:
+        ...
 
-TDataClass = TypeVar("TDataClass", bound=DataClass)
+    __dataclass_fields__: ClassVar[DataClassFields]
 
 
 Attr = Annotated[T, FieldType.ATTR]