Merge pull request #136 from astropenguin/astropenguin/issue135

Preferentially use data options if they are given in asdata* functions

Author: astropenguin

README.md

@@ -311,14 +311,11 @@ class Custom(xr.DataArray):
         print("Custom method!")
 
 
-dataoptions = DataOptions(Custom)
-
-
 @dataclass
 class Image(AsDataArray):
     """Specs for a monochromatic image."""
 
-    __dataoptions__ = dataoptions
+    __dataoptions__ = DataOptions(Custom)
 
     data: Data[tuple[X, Y], float]
     x: Coord[X, int] = 0

tests/test_dataarray.py

@@ -32,14 +32,11 @@ class Custom(xr.DataArray):
     __slots__ = ()
 
 
-dataoptions = DataOptions(Custom)
-
-
 @dataclass
 class Image(AsDataArray):
     """Specs for a monochromatic image."""
 
-    __dataoptions__ = dataoptions
+    __dataoptions__ = DataOptions(Custom)
 
     data: Data[Tuple[X, Y], float]
     x: Coord[X, int] = 0

tests/test_dataset.py

@@ -30,9 +30,6 @@ class Custom(xr.Dataset):
     __slots__ = ()
 
 
-dataoptions = DataOptions(Custom)
-
-
 @dataclass
 class Image(AsDataArray):
     """Specs for a monochromatic image."""
@@ -44,7 +41,7 @@ class Image(AsDataArray):
 class ColorImage(AsDataset):
     """Specs for a color image."""
 
-    __dataoptions__ = dataoptions
+    __dataoptions__ = DataOptions(Custom)
 
     red: Data[Tuple[X, Y], float]
     green: Data[Tuple[X, Y], float]

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,60 +26,65 @@
 # submodules
 from .datamodel import DataModel
 from .dataoptions import DataOptions
-from .typing import DataType, Order, Shape, Sizes
-
-
-# constants
-DEFAULT_OPTIONS = DataOptions(xr.DataArray)
+from .typing import DataClass, DataClassFields, DataType, Order, Shape, Sizes
 
 
 # type hints
 P = ParamSpec("P")
 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]]
+    __dataclass_fields__: ClassVar[DataClassFields]
+    __dataoptions__: DataOptions[TDataArray]
 
 
-class DataArrayClass(Protocol[P, TDataArray]):
-    """Type hint for a dataclass object with a DataArray factory."""
+# runtime functions
+@overload
+def asdataarray(
+    dataclass: OptionedClass[P, TDataArray],
+    reference: Optional[DataType] = None,
+    dataoptions: None = None,
+) -> TDataArray:
+    ...
 
-    def __init__(self, *args: P.args, **kwargs: P.kwargs) -> None:
-        ...
 
-    __dataclass_fields__: Dict[str, Field[Any]]
-    __dataoptions__: DataOptions[TDataArray]
+@overload
+def asdataarray(
+    dataclass: DataClass[P],
+    reference: Optional[DataType] = None,
+    dataoptions: None = None,
+) -> xr.DataArray:
+    ...
 
 
-# runtime functions
 @overload
 def asdataarray(
-    dataclass: DataArrayClass[Any, TDataArray],
+    dataclass: OptionedClass[P, Any],
     reference: Optional[DataType] = None,
-    dataoptions: DataOptions[Any] = DEFAULT_OPTIONS,
+    dataoptions: Optional[DataOptions[TDataArray]] = None,
 ) -> TDataArray:
     ...
 
 
 @overload
 def asdataarray(
-    dataclass: DataClass[Any],
+    dataclass: DataClass[P],
     reference: Optional[DataType] = None,
-    dataoptions: DataOptions[TDataArray] = DEFAULT_OPTIONS,
+    dataoptions: Optional[DataOptions[TDataArray]] = None,
 ) -> TDataArray:
     ...
 
 
 def asdataarray(
     dataclass: Any,
     reference: Optional[DataType] = None,
-    dataoptions: DataOptions[Any] = DEFAULT_OPTIONS,
+    dataoptions: Any = None,
 ) -> Any:
     """Create a DataArray object from a dataclass object.
 
@@ -84,10 +97,11 @@ def asdataarray(
         DataArray object created from the dataclass object.
 
     """
-    try:
-        dataoptions = dataclass.__dataoptions__
-    except AttributeError:
-        pass
+    if dataoptions is None:
+        try:
+            dataoptions = dataclass.__dataoptions__
+        except AttributeError:
+            dataoptions = DataOptions(xr.DataArray)
 
     model = DataModel.from_dataclass(dataclass)
     item = next(iter(model.data.values()))
@@ -127,7 +141,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 +177,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 +223,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 +265,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 +307,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 +349,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