Repository: astropenguin/xarray-dataclasses (python)

Issue:
Add factory fields for custom DataArray/Dataset creation
Add support of special fields (`__dataarray_factory__`, `__dataset_factory__`) for custom DataArray or Dataset creation.

```python
class CustomDataset(xr.Dataset):
    __slots__ = ()


@dataclass
class Custom(DataArrayMixin):
    data: Data[tuple["x", "y"], float]
    __dataset_factory__ = CustomDataset


ds = asdataset(Custom(...)) # statically typed as CustomDataset
type(ds) # -> CustomDataset
```

Pull request description:
Add factory fields for custom DataArray/Dataset creation
This PR closes #50.

Patch:
diff --git a/package-lock.json b/package-lock.json
index 16ba2d0..4e06741 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -5,13 +5,13 @@
   "packages": {
     "": {
       "devDependencies": {
-        "pyright": "^1.1.126"
+        "pyright": "~1.1"
       }
     },
     "node_modules/pyright": {
-      "version": "1.1.126",
-      "resolved": "https://registry.npmjs.org/pyright/-/pyright-1.1.126.tgz",
-      "integrity": "sha512-qOWOMM0YWsX7nZ224DjE5eQoJPqCCkzlkDAxSzCJESnM/Hs6SXO+Tq3VrsiePHSDxIzpDABlB6nQ7WLxd0f1gg==",
+      "version": "1.1.153",
+      "resolved": "https://registry.npmjs.org/pyright/-/pyright-1.1.153.tgz",
+      "integrity": "sha512-t5Y9IojL2Sifm1Q04BEBJDLP86rF4RpXma28HEP/EXJRHyyNxu8css77HC+tZqJXXXZVtAjaJxJSAEwq5afGjg==",
       "dev": true,
       "bin": {
         "pyright": "index.js",
@@ -24,9 +24,9 @@
   },
   "dependencies": {
     "pyright": {
-      "version": "1.1.126",
-      "resolved": "https://registry.npmjs.org/pyright/-/pyright-1.1.126.tgz",
-      "integrity": "sha512-qOWOMM0YWsX7nZ224DjE5eQoJPqCCkzlkDAxSzCJESnM/Hs6SXO+Tq3VrsiePHSDxIzpDABlB6nQ7WLxd0f1gg==",
+      "version": "1.1.153",
+      "resolved": "https://registry.npmjs.org/pyright/-/pyright-1.1.153.tgz",
+      "integrity": "sha512-t5Y9IojL2Sifm1Q04BEBJDLP86rF4RpXma28HEP/EXJRHyyNxu8css77HC+tZqJXXXZVtAjaJxJSAEwq5afGjg==",
       "dev": true
     }
   }
diff --git a/package.json b/package.json
index 92d322d..994e699 100644
--- a/package.json
+++ b/package.json
@@ -1,6 +1,6 @@
 {
   "devDependencies": {
-    "pyright": "^1.1.126"
+    "pyright": "~1.1"
   },
   "scripts": {
     "pyright": "poetry run pyright"
diff --git a/xarray_dataclasses/common.py b/xarray_dataclasses/common.py
index 81b7ac4..7142d3c 100644
--- a/xarray_dataclasses/common.py
+++ b/xarray_dataclasses/common.py
@@ -138,7 +138,7 @@ def _gen_fields(
     """
     for name, field in obj.__dataclass_fields__.items():
         if type_filter is None or type_filter(field.type):
-            yield field, getattr(obj, name, MISSING)
+            yield field, getattr(obj, name, MISSING)  # type: ignore
 
 
 def _to_dataarray(
diff --git a/xarray_dataclasses/dataarray.py b/xarray_dataclasses/dataarray.py
index 7c87bed..78fbe91 100644
--- a/xarray_dataclasses/dataarray.py
+++ b/xarray_dataclasses/dataarray.py
@@ -4,14 +4,13 @@ __all__ = ["asdataarray", "dataarrayclass"]
 # standard library
 from dataclasses import dataclass
 from functools import wraps
-from types import FunctionType
-from typing import Any, Callable, cast, Optional, Sequence, Type, Union
+from typing import Any, Callable, Optional, overload, Sequence, Type, TypeVar, Union
 
 
 # third-party packages
 import numpy as np
 import xarray as xr
-from typing_extensions import Literal
+from typing_extensions import Literal, Protocol
 
 
 # submodules
@@ -27,12 +26,38 @@ TEMP_CLASS_PREFIX: str = "__Copied"
 # type hints (internal)
 Order = Literal["C", "F"]
 Shape = Union[Sequence[int], int]
+DA = TypeVar("DA", covariant=True, bound=xr.DataArray)
+
+
+class DataClassWithFactory(DataClass, Protocol[DA]):
+    __dataarray_factory__: Callable[..., DA]
 
 
 # runtime functions (public)
-def asdataarray(inst: DataClass) -> xr.DataArray:
+@overload
+def asdataarray(
+    inst: DataClassWithFactory[DA],
+    dataarray_factory: Type[Any] = xr.DataArray,
+) -> DA:
+    ...
+
+
+@overload
+def asdataarray(
+    inst: DataClass,
+    dataarray_factory: Type[DA] = xr.DataArray,
+) -> DA:
+    ...
+
+
+def asdataarray(inst: Any, dataarray_factory: Any = xr.DataArray) -> Any:
     """Convert a DataArray-class instance to DataArray one."""
-    dataarray = get_data(inst)
+    try:
+        dataarray_factory = inst.__dataarray_factory__
+    except AttributeError:
+        pass
+
+    dataarray = dataarray_factory(get_data(inst))
     coords = get_coords(inst, dataarray)
 
     dataarray.coords.update(coords)
@@ -78,22 +103,24 @@ def dataarrayclass(
 class DataArrayMixin:
     """Mix-in class that provides shorthand methods."""
 
+    __dataarray_factory__ = xr.DataArray
+
     @classmethod
     def new(
-        cls,
+        cls: Type[DataClassWithFactory[DA]],
         *args: Any,
         **kwargs: Any,
-    ) -> xr.DataArray:
+    ) -> DA:
         """Create a DataArray instance."""
         raise NotImplementedError
 
     @classmethod
     def empty(
-        cls,
+        cls: Type[DataClassWithFactory[DA]],
         shape: Shape,
         order: Order = "C",
         **kwargs: Any,
-    ) -> xr.DataArray:
+    ) -> DA:
         """Create a DataArray instance without initializing data.
 
         Args:
@@ -106,18 +133,17 @@ class DataArrayMixin:
             DataArray instance filled without initializing data.
 
         """
-        cls = cast(Type[DataClass], cls)
         name = get_data_name(cls)
         data = np.empty(shape, order=order)
         return asdataarray(cls(**{name: data}, **kwargs))
 
     @classmethod
     def zeros(
-        cls,
+        cls: Type[DataClassWithFactory[DA]],
         shape: Shape,
         order: Order = "C",
         **kwargs: Any,
-    ) -> xr.DataArray:
+    ) -> DA:
         """Create a DataArray instance filled with zeros.
 
         Args:
@@ -130,18 +156,17 @@ class DataArrayMixin:
             DataArray instance filled with zeros.
 
         """
-        cls = cast(Type[DataClass], cls)
         name = get_data_name(cls)
         data = np.zeros(shape, order=order)
         return asdataarray(cls(**{name: data}, **kwargs))
 
     @classmethod
     def ones(
-        cls,
+        cls: Type[DataClassWithFactory[DA]],
         shape: Shape,
         order: Order = "C",
         **kwargs: Any,
-    ) -> xr.DataArray:
+    ) -> DA:
         """Create a DataArray instance filled with ones.
 
         Args:
@@ -154,19 +179,18 @@ class DataArrayMixin:
             DataArray instance filled with ones.
 
         """
-        cls = cast(Type[DataClass], cls)
         name = get_data_name(cls)
         data = np.ones(shape, order=order)
         return asdataarray(cls(**{name: data}, **kwargs))
 
     @classmethod
     def full(
-        cls,
+        cls: Type[DataClassWithFactory[DA]],
         shape: Shape,
         fill_value: Any,
         order: Order = "C",
         **kwargs: Any,
-    ) -> xr.DataArray:
+    ) -> DA:
         """Create a DataArray instance filled with given value.
 
         Args:
@@ -180,7 +204,6 @@ class DataArrayMixin:
             DataArray instance filled with given value.
 
         """
-        cls = cast(Type[DataClass], cls)
         name = get_data_name(cls)
         data = np.full(shape, fill_value, order=order)
         return asdataarray(cls(**{name: data}, **kwargs))
@@ -195,19 +218,17 @@ class DataArrayMixin:
         except ValueError:
             return
 
-        init: FunctionType = Temp.__init__  # type: ignore
-        init.__annotations__["return"] = xr.DataArray
+        init = Temp.__init__
+        init.__annotations__["return"] = DA
 
         # create a concrete new method and bind
         @classmethod
         @wraps(init)
         def new(
-            cls,  # type: ignore
+            cls: Type[DataClassWithFactory[DA]],
             *args: Any,
             **kwargs: Any,
-        ) -> xr.DataArray:
-            """Create a DataArray instance."""
-            cls = cast(Type[DataClass], cls)
+        ) -> DA:
             return asdataarray(cls(*args, **kwargs))
 
         cls.new = new  # type: ignore
diff --git a/xarray_dataclasses/dataset.py b/xarray_dataclasses/dataset.py
index 388d7f8..0ede43e 100644
--- a/xarray_dataclasses/dataset.py
+++ b/xarray_dataclasses/dataset.py
@@ -4,12 +4,12 @@ __all__ = ["asdataset", "datasetclass"]
 # standard library
 from dataclasses import dataclass
 from functools import wraps
-from types import FunctionType
-from typing import Any, Callable, cast, Optional, Type, Union
+from typing import Any, Callable, Optional, overload, Type, TypeVar, Union
 
 
 # third-party packages
 import xarray as xr
+from typing_extensions import Protocol
 
 
 # submodules
@@ -22,10 +22,39 @@ from .utils import copy_class, extend_class
 TEMP_CLASS_PREFIX: str = "__Copied"
 
 
+# type hints (internal)
+DS = TypeVar("DS", covariant=True, bound=xr.Dataset)
+
+
+class DataClassWithFactory(DataClass, Protocol[DS]):
+    __dataset_factory__: Callable[..., DS]
+
+
 # runtime functions (public)
-def asdataset(inst: DataClass) -> xr.Dataset:
+@overload
+def asdataset(
+    inst: DataClassWithFactory[DS],
+    dataset_factory: Type[Any] = xr.Dataset,
+) -> DS:
+    ...
+
+
+@overload
+def asdataset(
+    inst: DataClass,
+    dataset_factory: Type[DS] = xr.Dataset,
+) -> DS:
+    ...
+
+
+def asdataset(inst: Any, dataset_factory: Any = xr.Dataset) -> Any:
     """Convert a Dataset-class instance to Dataset one."""
-    dataset = xr.Dataset(get_data_vars(inst))
+    try:
+        dataset_factory = inst.__dataset_factory__
+    except AttributeError:
+        pass
+
+    dataset = dataset_factory(get_data_vars(inst))
     coords = get_coords(inst, dataset)
 
     dataset.coords.update(coords)
@@ -70,12 +99,14 @@ def datasetclass(
 class DatasetMixin:
     """Mix-in class that provides shorthand methods."""
 
+    __dataset_factory__ = xr.Dataset
+
     @classmethod
     def new(
-        cls,
+        cls: Type[DataClassWithFactory[DS]],
         *args: Any,
         **kwargs: Any,
-    ) -> xr.Dataset:
+    ) -> DS:
         """Create a Dataset instance."""
         raise NotImplementedError
 
@@ -89,19 +120,17 @@ class DatasetMixin:
         except ValueError:
             return
 
-        init: FunctionType = Temp.__init__  # type: ignore
-        init.__annotations__["return"] = xr.Dataset
+        init = Temp.__init__
+        init.__annotations__["return"] = DS
 
         # create a concrete new method and bind
         @classmethod
         @wraps(init)
         def new(
-            cls,  # type: ignore
+            cls: Type[DataClassWithFactory[DS]],
             *args: Any,
             **kwargs: Any,
-        ) -> xr.Dataset:
-            """Create a Dataset instance."""
-            cls = cast(Type[DataClass], cls)
+        ) -> DS:
             return asdataset(cls(*args, **kwargs))
 
         cls.new = new  # type: ignore


Test patch:
diff --git a/tests/test_dataarray.py b/tests/test_dataarray.py
index 577086c..d386f07 100644
--- a/tests/test_dataarray.py
+++ b/tests/test_dataarray.py
@@ -25,8 +25,14 @@ Y = Literal[DIMS[1]]
 
 
 # dataclasses
+class Custom(xr.DataArray):
+    __slots__ = ()
+
+
 @dataclass
 class Image(DataArrayMixin):
+    __dataarray_factory__ = Custom
+
     data: Data[Tuple[X, Y], float]
     x: Coord[X, int] = 0
     y: Coord[Y, int] = 0
@@ -36,7 +42,7 @@ class Image(DataArrayMixin):
 
 # test datasets
 created = Image.ones(SHAPE)
-expected = xr.DataArray(
+expected = Custom(
     np.ones(SHAPE, float),
     dims=DIMS,
     coords={
@@ -49,6 +55,10 @@ expected = xr.DataArray(
 
 
 # test functions
+def test_type() -> None:
+    assert type(created) is type(expected)
+
+
 def test_data() -> None:
     assert (created == expected).all()  # type: ignore
 
diff --git a/tests/test_dataset.py b/tests/test_dataset.py
index 45d6b48..c5be22f 100644
--- a/tests/test_dataset.py
+++ b/tests/test_dataset.py
@@ -25,6 +25,10 @@ Y = Literal[DIMS[1]]
 
 
 # dataclasses
+class Custom(xr.Dataset):
+    __slots__ = ()
+
+
 @dataclass
 class Image(DataArrayMixin):
     data: Data[Tuple[X, Y], float]
@@ -32,6 +36,8 @@ class Image(DataArrayMixin):
 
 @dataclass
 class RGBImage(DatasetMixin):
+    __dataset_factory__ = Custom
+
     red: Data[Tuple[X, Y], float]
     green: Data[Tuple[X, Y], float]
     blue: Data[Tuple[X, Y], float]
@@ -46,7 +52,7 @@ created = RGBImage.new(
     Image.ones(SHAPE),
     Image.ones(SHAPE),
 )
-expected = xr.Dataset(
+expected = Custom(
     data_vars={
         "red": xr.DataArray(np.ones(SHAPE), dims=DIMS),
         "green": xr.DataArray(np.ones(SHAPE), dims=DIMS),
@@ -61,6 +67,10 @@ expected = xr.Dataset(
 
 
 # test functions
+def test_type() -> None:
+    assert type(created) is type(expected)
+
+
 def test_data_vars() -> None:
     assert (created == expected).all()  # type: ignore