Merge pull request #124 from astropenguin:astropenguin/issue123

Add classmethod to create a shaped array

Author: astropenguin

pyrightconfig.json

@@ -1,6 +1,8 @@
 {
     "typeCheckingMode": "strict",
     "reportImportCycles": "warning",
+    "reportUnknownArgumentType": "warning",
     "reportUnknownMemberType": "warning",
+    "reportUnknownParameterType": "warning",
     "reportUnknownVariableType": "warning"
 }

xarray_dataclasses/dataarray.py

@@ -3,7 +3,7 @@
 
 # standard library
 from dataclasses import Field
-from functools import wraps
+from functools import partial, wraps
 from types import MethodType
 from typing import Any, Callable, Dict, Optional, Type, TypeVar, Union, overload
 
@@ -166,6 +166,52 @@ def new(cls: Any, *args: Any, **kwargs: Any) -> Any:
 
         return MethodType(new, cls)
 
+    @overload
+    @classmethod
+    def shaped(
+        cls: Type[DataArrayClass[P, TDataArray]],
+        func: Callable[[Shape], np.ndarray],
+        shape: Union[Shape, Sizes],
+        **kwargs: Any,
+    ) -> TDataArray:
+        ...
+
+    @overload
+    @classmethod
+    def shaped(
+        cls: Type[DataClass[P]],
+        func: Callable[[Shape], np.ndarray],
+        shape: Union[Shape, Sizes],
+        **kwargs: Any,
+    ) -> xr.DataArray:
+        ...
+
+    @classmethod
+    def shaped(
+        cls: Any,
+        func: Callable[[Shape], np.ndarray],
+        shape: Union[Shape, Sizes],
+        **kwargs: Any,
+    ) -> Any:
+        """Create a DataArray object from a shaped function.
+
+        Args:
+            func: Function to create an array with given shape.
+            shape: Shape or sizes of the new DataArray object.
+            kwargs: Args of the DataArray class except for data.
+
+        Returns:
+            DataArray object created from the shaped function.
+
+        """
+        model = DataModel.from_dataclass(cls)
+        name, item = next(iter(model.data.items()))
+
+        if isinstance(shape, dict):
+            shape = tuple(shape[dim] for dim in item.type["dims"])
+
+        return asdataarray(cls(**{name: func(shape)}, **kwargs))
+
     @overload
     @classmethod
     def empty(
@@ -205,14 +251,8 @@ def empty(
             DataArray object without initializing data.
 
         """
-        model = DataModel.from_dataclass(cls)
-        name, item = next(iter(model.data.items()))
-
-        if isinstance(shape, dict):
-            shape = tuple(shape[dim] for dim in item.type["dims"])
-
-        data = np.empty(shape, order=order)
-        return asdataarray(cls(**{name: data}, **kwargs))
+        func = partial(np.empty, order=order)
+        return cls.shaped(func, shape, **kwargs)
 
     @overload
     @classmethod
@@ -253,14 +293,8 @@ def zeros(
             DataArray object filled with zeros.
 
         """
-        model = DataModel.from_dataclass(cls)
-        name, item = next(iter(model.data.items()))
-
-        if isinstance(shape, dict):
-            shape = tuple(shape[dim] for dim in item.type["dims"])
-
-        data = np.zeros(shape, order=order)
-        return asdataarray(cls(**{name: data}, **kwargs))
+        func = partial(np.zeros, order=order)
+        return cls.shaped(func, shape, **kwargs)
 
     @overload
     @classmethod
@@ -301,14 +335,8 @@ def ones(
             DataArray object filled with ones.
 
         """
-        model = DataModel.from_dataclass(cls)
-        name, item = next(iter(model.data.items()))
-
-        if isinstance(shape, dict):
-            shape = tuple(shape[dim] for dim in item.type["dims"])
-
-        data = np.ones(shape, order=order)
-        return asdataarray(cls(**{name: data}, **kwargs))
+        func = partial(np.ones, order=order)
+        return cls.shaped(func, shape, **kwargs)
 
     @overload
     @classmethod
@@ -353,11 +381,5 @@ def full(
             DataArray object filled with given value.
 
         """
-        model = DataModel.from_dataclass(cls)
-        name, item = next(iter(model.data.items()))
-
-        if isinstance(shape, dict):
-            shape = tuple(shape[dim] for dim in item.type["dims"])
-
-        data = np.full(shape, fill_value, order=order)
-        return asdataarray(cls(**{name: data}, **kwargs))
+        func = partial(np.full, fill_value=fill_value, order=order)
+        return cls.shaped(func, shape, **kwargs)

xarray_dataclasses/dataset.py

@@ -3,7 +3,7 @@
 
 # standard library
 from dataclasses import Field
-from functools import wraps
+from functools import partial, wraps
 from types import MethodType
 from typing import Any, Callable, Dict, Optional, Type, TypeVar, overload
 
@@ -18,7 +18,7 @@
 # submodules
 from .datamodel import DataModel
 from .dataoptions import DataOptions
-from .typing import DataType, Order, Sizes
+from .typing import DataType, Order, Shape, Sizes
 
 
 # constants
@@ -164,6 +164,53 @@ def new(cls: Any, *args: Any, **kwargs: Any) -> Any:
 
         return MethodType(new, cls)
 
+    @overload
+    @classmethod
+    def shaped(
+        cls: Type[DatasetClass[P, TDataset]],
+        func: Callable[[Shape], np.ndarray],
+        sizes: Sizes,
+        **kwargs: Any,
+    ) -> TDataset:
+        ...
+
+    @overload
+    @classmethod
+    def shaped(
+        cls: Type[DataClass[P]],
+        func: Callable[[Shape], np.ndarray],
+        sizes: Sizes,
+        **kwargs: Any,
+    ) -> xr.Dataset:
+        ...
+
+    @classmethod
+    def shaped(
+        cls: Any,
+        func: Callable[[Shape], np.ndarray],
+        sizes: Sizes,
+        **kwargs: Any,
+    ) -> Any:
+        """Create a Dataset object from a shaped function.
+
+        Args:
+            func: Function to create an array with given shape.
+            sizes: Sizes of the new Dataset object.
+            kwargs: Args of the Dataset class except for data vars.
+
+        Returns:
+            Dataset object created from the shaped function.
+
+        """
+        model = DataModel.from_dataclass(cls)
+        data_vars: Dict[str, Any] = {}
+
+        for name, item in model.data.items():
+            shape = tuple(sizes[dim] for dim in item.type["dims"])
+            data_vars[name] = func(shape)
+
+        return asdataset(cls(**data_vars, **kwargs))
+
     @overload
     @classmethod
     def empty(
@@ -203,14 +250,8 @@ def empty(
             Dataset object without initializing data vars.
 
         """
-        model = DataModel.from_dataclass(cls)
-        data_vars: Dict[str, Any] = {}
-
-        for name, item in model.data.items():
-            shape = tuple(sizes[dim] for dim in item.type["dims"])
-            data_vars[name] = np.empty(shape, order=order)
-
-        return asdataset(cls(**data_vars, **kwargs))
+        func = partial(np.empty, order=order)
+        return cls.shaped(func, sizes, **kwargs)
 
     @overload
     @classmethod
@@ -251,14 +292,8 @@ def zeros(
             Dataset object whose data vars are filled with zeros.
 
         """
-        model = DataModel.from_dataclass(cls)
-        data_vars: Dict[str, Any] = {}
-
-        for name, item in model.data.items():
-            shape = tuple(sizes[dim] for dim in item.type["dims"])
-            data_vars[name] = np.zeros(shape, order=order)
-
-        return asdataset(cls(**data_vars, **kwargs))
+        func = partial(np.zeros, order=order)
+        return cls.shaped(func, sizes, **kwargs)
 
     @overload
     @classmethod
@@ -299,14 +334,8 @@ def ones(
             Dataset object whose data vars are filled with ones.
 
         """
-        model = DataModel.from_dataclass(cls)
-        data_vars: Dict[str, Any] = {}
-
-        for name, item in model.data.items():
-            shape = tuple(sizes[dim] for dim in item.type["dims"])
-            data_vars[name] = np.ones(shape, order=order)
-
-        return asdataset(cls(**data_vars, **kwargs))
+        func = partial(np.ones, order=order)
+        return cls.shaped(func, sizes, **kwargs)
 
     @overload
     @classmethod
@@ -351,11 +380,5 @@ def full(
             Dataset object whose data vars are filled with given value.
 
         """
-        model = DataModel.from_dataclass(cls)
-        data_vars: Dict[str, Any] = {}
-
-        for name, item in model.data.items():
-            shape = tuple(sizes[dim] for dim in item.type["dims"])
-            data_vars[name] = np.full(shape, fill_value, order=order)
-
-        return asdataset(cls(**data_vars, **kwargs))
+        func = partial(np.full, fill_value=fill_value, order=order)
+        return cls.shaped(func, sizes, **kwargs)