Repository: astropenguin/xarray-dataclasses

Issue:
Add data options
Add `__dataoptions__` that allows users to control detailed behaviors of DataArray or Dataset creation.

Comments:

Patch:
diff --git a/README.md b/README.md
index 7f5efcc..cfd55fd 100644
--- a/README.md
+++ b/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
diff --git a/xarray_dataclasses/__init__.py b/xarray_dataclasses/__init__.py
index 398052a..11f3de4 100644
--- a/xarray_dataclasses/__init__.py
+++ b/xarray_dataclasses/__init__.py
@@ -20,6 +20,7 @@ from . import dataarray
 from . import dataset
 from . import deprecated
 from . import datamodel
+from . import dataoptions
 from . import typing
 
 
@@ -28,6 +29,7 @@ from .dataarray import *
 from .dataset import *
 from .deprecated import *
 from .datamodel import *
+from .dataoptions import *
 from .typing import *
 
 
diff --git a/xarray_dataclasses/dataarray.py b/xarray_dataclasses/dataarray.py
index 10a99ba..7d6afa5 100644
--- a/xarray_dataclasses/dataarray.py
+++ b/xarray_dataclasses/dataarray.py
@@ -5,7 +5,7 @@ __all__ = ["AsDataArray", "asdataarray"]
 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 @@ from typing_extensions import ParamSpec, Protocol
 
 
 # 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,8 +70,8 @@ class classproperty:
 @overload
 def asdataarray(
     dataclass: DataArrayClass[Any, TDataArray],
-    reference: Reference = None,
-    dataarray_factory: Any = xr.DataArray,
+    reference: Optional[DataType] = None,
+    dataoptions: Any = DEFAULT_OPTIONS,
 ) -> TDataArray:
     ...
 
@@ -74,8 +79,8 @@ def asdataarray(
 @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:
     ...
 
@@ -83,26 +88,32 @@ def asdataarray(
 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]:
diff --git a/xarray_dataclasses/datamodel.py b/xarray_dataclasses/datamodel.py
index a865899..2d3e3c6 100644
--- a/xarray_dataclasses/datamodel.py
+++ b/xarray_dataclasses/datamodel.py
@@ -3,7 +3,7 @@ __all__ = ["DataModel"]
 
 # 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 .deprecated import get_type_hints
 from .typing import (
     ArrayLike,
     DataClass,
+    DataType,
     Dims,
     Dtype,
     FieldType,
@@ -28,8 +29,7 @@ from .typing import (
 
 
 # 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.
 
diff --git a/xarray_dataclasses/dataoptions.py b/xarray_dataclasses/dataoptions.py
new file mode 100644
index 0000000..3d479a3
--- /dev/null
+++ b/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."""
diff --git a/xarray_dataclasses/dataset.py b/xarray_dataclasses/dataset.py
index f71a1c9..507a954 100644
--- a/xarray_dataclasses/dataset.py
+++ b/xarray_dataclasses/dataset.py
@@ -5,7 +5,7 @@ __all__ = ["AsDataset", "asdataset"]
 from dataclasses import Field
 from functools import wraps
 from types import MethodType
-from typing import Any, Callable, Dict, Type, TypeVar, overload
+from typing import Any, Callable, Dict, Optional, Type, TypeVar, overload
 
 
 # dependencies
@@ -16,8 +16,13 @@ from typing_extensions import ParamSpec, Protocol
 
 
 # submodules
-from .datamodel import DataModel, Reference
-from .typing import Order, Sizes
+from .datamodel import DataModel
+from .dataoptions import DataOptions
+from .typing import DataType, Order, Sizes
+
+
+# constants
+DEFAULT_OPTIONS = DataOptions(xr.Dataset)
 
 
 # type hints
@@ -38,7 +43,7 @@ class DatasetClass(Protocol[P, TDataset_]):
 
     __init__: Callable[P, None]
     __dataclass_fields__: Dict[str, Field[Any]]
-    __dataset_factory__: Callable[..., TDataset_]
+    __dataoptions__: DataOptions[TDataset_]
 
 
 # custom classproperty
@@ -65,8 +70,8 @@ class classproperty:
 @overload
 def asdataset(
     dataclass: DatasetClass[Any, TDataset],
-    reference: Reference = None,
-    dataset_factory: Any = xr.Dataset,
+    reference: Optional[DataType] = None,
+    dataoptions: Any = DEFAULT_OPTIONS,
 ) -> TDataset:
     ...
 
@@ -74,8 +79,8 @@ def asdataset(
 @overload
 def asdataset(
     dataclass: DataClass[Any],
-    reference: Reference = None,
-    dataset_factory: Callable[..., TDataset] = xr.Dataset,
+    reference: Optional[DataType] = None,
+    dataoptions: DataOptions[TDataset] = DEFAULT_OPTIONS,
 ) -> TDataset:
     ...
 
@@ -83,26 +88,32 @@ def asdataset(
 def asdataset(
     dataclass: Any,
     reference: Any = None,
-    dataset_factory: Any = xr.Dataset,
+    dataoptions: Any = DEFAULT_OPTIONS,
 ) -> Any:
     """Create a Dataset object from a dataclass object.
 
     Args:
         dataclass: Dataclass object that defines typed Dataset.
         reference: DataArray or Dataset object as a reference of shape.
-        dataset_factory: Factory function of Dataset.
+        dataoptions: Options for Dataset creation.
 
     Returns:
         Dataset object created from the dataclass object.
 
     """
     try:
-        dataset_factory = dataclass.__dataset_factory__
+        # for backward compatibility (deprecated in v1.0.0)
+        dataoptions = DataOptions(dataclass.__dataset_factory__)
+    except AttributeError:
+        pass
+
+    try:
+        dataoptions = dataclass.__dataoptions__
     except AttributeError:
         pass
 
     model = DataModel.from_dataclass(dataclass)
-    dataset = dataset_factory()
+    dataset = dataoptions.factory()
 
     for data in model.data:
         dataset.update({data.name: data(reference)})
@@ -119,9 +130,7 @@ def asdataset(
 class AsDataset:
     """Mix-in class that provides shorthand methods."""
 
-    def __dataset_factory__(self, data_vars: Any = None) -> xr.Dataset:
-        """Default Dataset factory (xarray.Dataset)."""
-        return xr.Dataset(data_vars)
+    __dataoptions__ = DEFAULT_OPTIONS
 
     @classproperty
     def new(cls: Type[DatasetClass[P, TDataset]]) -> Callable[P, TDataset]:
diff --git a/xarray_dataclasses/typing.py b/xarray_dataclasses/typing.py
index a6c0043..fc67d01 100644
--- a/xarray_dataclasses/typing.py
+++ b/xarray_dataclasses/typing.py
@@ -33,6 +33,7 @@ from typing import (
 
 
 # dependencies
+import xarray as xr
 from typing_extensions import (
     Annotated,
     Literal,
@@ -72,6 +73,7 @@ class FieldType(Enum):
 
 
 # type hints
+DataType = Union[xr.DataArray, xr.Dataset]
 Dims = Tuple[str, ...]
 Dtype = Optional[str]
 Order = Literal["C", "F"]


Test patch:
diff --git a/tests/test_dataarray.py b/tests/test_dataarray.py
index aa77377..88f982b 100644
--- a/tests/test_dataarray.py
+++ b/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 @@ from typing_extensions import Literal
 
 # 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)
diff --git a/tests/test_dataset.py b/tests/test_dataset.py
index 35c4e1d..6169033 100644
--- a/tests/test_dataset.py
+++ b/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 @@ from typing_extensions import Literal
 # 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,6 +44,8 @@ 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]
@@ -47,9 +53,6 @@ class ColorImage(AsDataset):
     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(