Merge pull request #113 from astropenguin/astropenguin/issue112

Add data options

Author: astropenguin

README.md

@@ -291,12 +291,15 @@ class Image(AsDataArray):
     y: Coordof[YAxis] = 0
 ```
 
-### Custom DataArray and Dataset factories
+### Options for DataArray and Dataset creation
+
+For customization, users can add a special class attribute, `__dataoptions__`, to a DataArray or Dataset class.
+A custom factory for DataArray or Dataset creation is only supported in the current implementation.
 
-For customization, users can use a function or a class to create an initial DataArray or Dataset object by specifying a special class attribute, `__dataarray_factory__` or `__dataset_factory__`, respectively.
 
 ```python
 import xarray as xr
+from xarray_dataclasses import DataOptions
 
 
 class Custom(xr.DataArray):
@@ -308,19 +311,23 @@ class Custom(xr.DataArray):
         print("Custom method!")
 
 
+dataoptions = DataOptions(Custom)
+
+
 @dataclass
 class Image(AsDataArray):
     """Specs for a monochromatic image."""
 
+    __dataoptions__ = dataoptions
+
     data: Data[tuple[X, Y], float]
     x: Coord[X, int] = 0
     y: Coord[Y, int] = 0
-    __dataarray_factory__ = Custom
 
 
 image = Image.ones([3, 3])
-isinstance(image, Custom) # True
-image.custom_method() # Custom method!
+isinstance(image, Custom)  # True
+image.custom_method()  # Custom method!
 ```
 
 ### DataArray and Dataset creation without shorthands

tests/test_dataarray.py

@@ -1,6 +1,6 @@
 # standard library
 from dataclasses import dataclass
-from typing import Any, Tuple
+from typing import Tuple
 
 
 # third-party packages
@@ -11,6 +11,7 @@
 
 # submodules
 from xarray_dataclasses.dataarray import AsDataArray
+from xarray_dataclasses.dataoptions import DataOptions
 from xarray_dataclasses.typing import Attr, Coord, Data, Name
 
 
@@ -31,19 +32,21 @@ class Custom(xr.DataArray):
     __slots__ = ()
 
 
+dataoptions = DataOptions(Custom)
+
+
 @dataclass
 class Image(AsDataArray):
     """Specs for a monochromatic image."""
 
+    __dataoptions__ = dataoptions
+
     data: Data[Tuple[X, Y], float]
     x: Coord[X, int] = 0
     y: Coord[Y, int] = 0
     units: Attr[str] = "cd / m^2"
     name: Name[str] = "luminance"
 
-    def __dataarray_factory__(self, data: Any = None) -> Custom:
-        return Custom(data)
-
 
 # test datasets
 created = Image.ones(SHAPE)

tests/test_dataset.py

@@ -1,6 +1,6 @@
 # standard library
 from dataclasses import dataclass
-from typing import Any, Tuple
+from typing import Tuple
 
 
 # third-party packages
@@ -12,6 +12,7 @@
 # submodules
 from xarray_dataclasses.dataarray import AsDataArray
 from xarray_dataclasses.dataset import AsDataset
+from xarray_dataclasses.dataoptions import DataOptions
 from xarray_dataclasses.typing import Attr, Coord, Data
 
 # constants
@@ -29,6 +30,9 @@ class Custom(xr.Dataset):
     __slots__ = ()
 
 
+dataoptions = DataOptions(Custom)
+
+
 @dataclass
 class Image(AsDataArray):
     """Specs for a monochromatic image."""
@@ -40,16 +44,15 @@ class Image(AsDataArray):
 class ColorImage(AsDataset):
     """Specs for a color image."""
 
+    __dataoptions__ = dataoptions
+
     red: Data[Tuple[X, Y], float]
     green: Data[Tuple[X, Y], float]
     blue: Data[Tuple[X, Y], float]
     x: Coord[X, int] = 0
     y: Coord[Y, int] = 0
     units: Attr[str] = "cd / m^2"
 
-    def __dataset_factory__(self, data_vars: Any = None) -> Custom:
-        return Custom(data_vars)
-
 
 # test datasets
 created = ColorImage.new(

xarray_dataclasses/__init__.py

@@ -20,6 +20,7 @@ def _make_field_generic():
 from . import dataset
 from . import deprecated
 from . import datamodel
+from . import dataoptions
 from . import typing
 
 
@@ -28,6 +29,7 @@ def _make_field_generic():
 from .dataset import *
 from .deprecated import *
 from .datamodel import *
+from .dataoptions import *
 from .typing import *
 
 

xarray_dataclasses/dataarray.py

@@ -5,7 +5,7 @@
 from dataclasses import Field
 from functools import wraps
 from types import MethodType
-from typing import Any, Callable, Dict, Type, TypeVar, Union, overload
+from typing import Any, Callable, Dict, Optional, Type, TypeVar, Union, overload
 
 
 # dependencies
@@ -16,8 +16,13 @@
 
 
 # submodules
-from .datamodel import DataModel, Reference
-from .typing import Order, Shape, Sizes
+from .datamodel import DataModel
+from .dataoptions import DataOptions
+from .typing import DataType, Order, Shape, Sizes
+
+
+# constants
+DEFAULT_OPTIONS = DataOptions(xr.DataArray)
 
 
 # type hints
@@ -38,7 +43,7 @@ class DataArrayClass(Protocol[P, TDataArray_]):
 
     __init__: Callable[P, None]
     __dataclass_fields__: Dict[str, Field[Any]]
-    __dataarray_factory__: Callable[..., TDataArray_]
+    __dataoptions__: DataOptions[TDataArray_]
 
 
 # custom classproperty
@@ -65,44 +70,50 @@ def __get__(
 @overload
 def asdataarray(
     dataclass: DataArrayClass[Any, TDataArray],
-    reference: Reference = None,
-    dataarray_factory: Any = xr.DataArray,
+    reference: Optional[DataType] = None,
+    dataoptions: Any = DEFAULT_OPTIONS,
 ) -> TDataArray:
     ...
 
 
 @overload
 def asdataarray(
     dataclass: DataClass[Any],
-    reference: Reference = None,
-    dataarray_factory: Callable[..., TDataArray] = xr.DataArray,
+    reference: Optional[DataType] = None,
+    dataoptions: DataOptions[TDataArray] = DEFAULT_OPTIONS,
 ) -> TDataArray:
     ...
 
 
 def asdataarray(
     dataclass: Any,
     reference: Any = None,
-    dataarray_factory: Any = xr.DataArray,
+    dataoptions: Any = DEFAULT_OPTIONS,
 ) -> Any:
     """Create a DataArray object from a dataclass object.
 
     Args:
         dataclass: Dataclass object that defines typed DataArray.
         reference: DataArray or Dataset object as a reference of shape.
-        dataset_factory: Factory function of DataArray.
+        dataoptions: Options for DataArray creation.
 
     Returns:
         DataArray object created from the dataclass object.
 
     """
     try:
-        dataarray_factory = dataclass.__dataarray_factory__
+        # for backward compatibility (deprecated in v1.0.0)
+        dataoptions = DataOptions(dataclass.__dataarray_factory__)
+    except AttributeError:
+        pass
+
+    try:
+        dataoptions = dataclass.__dataoptions__
     except AttributeError:
         pass
 
     model = DataModel.from_dataclass(dataclass)
-    dataarray = dataarray_factory(model.data[0](reference))
+    dataarray = dataoptions.factory(model.data[0](reference))
 
     for coord in model.coord:
         dataarray.coords.update({coord.name: coord(dataarray)})
@@ -119,9 +130,7 @@ def asdataarray(
 class AsDataArray:
     """Mix-in class that provides shorthand methods."""
 
-    def __dataarray_factory__(self, data: Any = None) -> xr.DataArray:
-        """Default DataArray factory (xarray.DataArray)."""
-        return xr.DataArray(data)
+    __dataoptions__ = DEFAULT_OPTIONS
 
     @classproperty
     def new(cls: Type[DataArrayClass[P, TDataArray]]) -> Callable[P, TDataArray]:

xarray_dataclasses/datamodel.py

@@ -3,7 +3,7 @@
 
 # standard library
 from dataclasses import Field, dataclass, field, is_dataclass
-from typing import Any, List, Optional, Type, Union, cast
+from typing import Any, List, Optional, Type, cast
 
 
 # dependencies
@@ -17,6 +17,7 @@
 from .typing import (
     ArrayLike,
     DataClass,
+    DataType,
     Dims,
     Dtype,
     FieldType,
@@ -28,8 +29,7 @@
 
 
 # type hints
-DataType = TypedDict("DataType", dims=Dims, dtype=Dtype)
-Reference = Union[xr.DataArray, xr.Dataset, None]
+DimsDtype = TypedDict("DimsDtype", dims=Dims, dtype=Dtype)
 
 
 # field models
@@ -43,13 +43,13 @@ class Data:
     value: Any
     """Value assigned to the field."""
 
-    type: DataType
+    type: DimsDtype
     """Type (dims and dtype) of the field."""
 
     factory: Optional[Type[DataClass]] = None
     """Factory dataclass to create a DataArray object."""
 
-    def __call__(self, reference: Reference = None) -> xr.DataArray:
+    def __call__(self, reference: Optional[DataType] = None) -> xr.DataArray:
         """Create a DataArray object from the value and a reference."""
         from .dataarray import asdataarray
 
@@ -173,7 +173,7 @@ def typedarray(
     data: Any,
     dims: Dims,
     dtype: Dtype,
-    reference: Reference = None,
+    reference: Optional[DataType] = None,
 ) -> xr.DataArray:
     """Create a DataArray object with given dims and dtype.
 

xarray_dataclasses/dataoptions.py

@@ -0,0 +1,23 @@
+__all__ = ["DataOptions"]
+
+
+# standard library
+from dataclasses import dataclass
+from typing import Callable, Generic, TypeVar
+
+
+# submodules
+from .typing import DataType
+
+
+# type hints
+TDataType = TypeVar("TDataType", bound=DataType)
+
+
+# dataclasses
+@dataclass(frozen=True)
+class DataOptions(Generic[TDataType]):
+    """Options for DataArray or Dataset creation."""
+
+    factory: Callable[..., TDataType]
+    """Factory function for DataArray or Dataset."""