#121 Use overload in AsDataArray

Author: astropenguin

xarray_dataclasses/dataarray.py

@@ -49,32 +49,12 @@ def __init__(self, *args: P.args, **kwargs: P.kwargs) -> None:
     __dataoptions__: DataOptions[TDataArray]
 
 
-# custom classproperty
-class classproperty:
-    """Class property only for AsDataArray.new().
-
-    As a classmethod and a property can be chained together since Python 3.9,
-    this class will be removed when the support for Python 3.7 and 3.8 ends.
-
-    """
-
-    def __init__(self, func: Callable[..., Any]) -> None:
-        self.__func__ = func
-
-    def __get__(
-        self,
-        obj: Any,
-        cls: Type[DataArrayClass[P, TDataArray]],
-    ) -> Callable[P, TDataArray]:
-        return self.__func__(cls)
-
-
-# runtime functions and classes
+# runtime functions
 @overload
 def asdataarray(
     dataclass: DataArrayClass[Any, TDataArray],
     reference: Optional[DataType] = None,
-    dataoptions: Any = DEFAULT_OPTIONS,
+    dataoptions: DataOptions[Any] = DEFAULT_OPTIONS,
 ) -> TDataArray:
     ...
 
@@ -90,8 +70,8 @@ def asdataarray(
 
 def asdataarray(
     dataclass: Any,
-    reference: Any = None,
-    dataoptions: Any = DEFAULT_OPTIONS,
+    reference: Optional[DataType] = None,
+    dataoptions: DataOptions[Any] = DEFAULT_OPTIONS,
 ) -> Any:
     """Create a DataArray object from a dataclass object.
 
@@ -137,36 +117,82 @@ def asdataarray(
     return dataarray
 
 
+# runtime classes
+class classproperty:
+    """Class property only for AsDataArray.new().
+
+    As a classmethod and a property can be chained together since Python 3.9,
+    this class will be removed when the support for Python 3.7 and 3.8 ends.
+
+    """
+
+    def __init__(self, func: Any) -> None:
+        self.__func__ = func
+
+    @overload
+    def __get__(
+        self,
+        obj: Any,
+        cls: Type[DataArrayClass[P, TDataArray]],
+    ) -> Callable[P, TDataArray]:
+        ...
+
+    @overload
+    def __get__(
+        self,
+        obj: Any,
+        cls: Type[DataClass[P]],
+    ) -> Callable[P, xr.DataArray]:
+        ...
+
+    def __get__(self, obj: Any, cls: Any) -> Any:
+        return self.__func__(cls)
+
+
 class AsDataArray:
     """Mix-in class that provides shorthand methods."""
 
-    __dataoptions__ = DEFAULT_OPTIONS
-
     @classproperty
-    def new(cls: Type[DataArrayClass[P, TDataArray]]) -> Callable[P, TDataArray]:
+    def new(cls: Any) -> Any:
         """Create a DataArray object from dataclass parameters."""
 
-        init = copy(cls.__init__)  # type: ignore
-        init.__doc__ = cls.__init__.__doc__  # type: ignore
+        init = copy(cls.__init__)
+        init.__doc__ = cls.__init__.__doc__
         init.__annotations__["return"] = TDataArray
 
         @wraps(init)
-        def new(
-            cls: Type[DataArrayClass[P, TDataArray]],
-            *args: P.args,
-            **kwargs: P.kwargs,
-        ) -> TDataArray:
+        def new(cls: Any, *args: Any, **kwargs: Any) -> Any:
             return asdataarray(cls(*args, **kwargs))
 
         return MethodType(new, cls)
 
+    @overload
     @classmethod
     def empty(
         cls: Type[DataArrayClass[P, TDataArray]],
         shape: Union[Shape, Sizes],
         order: Order = "C",
         **kwargs: Any,
     ) -> TDataArray:
+        ...
+
+    @overload
+    @classmethod
+    def empty(
+        cls: Type[DataClass[P]],
+        shape: Union[Shape, Sizes],
+        order: Order = "C",
+        **kwargs: Any,
+    ) -> xr.DataArray:
+        ...
+
+    @classmethod
+    def empty(
+        cls: Any,
+        shape: Union[Shape, Sizes],
+        order: Order = "C",
+        **kwargs: Any,
+    ) -> Any:
         """Create a DataArray object without initializing data.
 
         Args:
@@ -188,13 +214,33 @@ def empty(
         data = np.empty(shape, order=order)
         return asdataarray(cls(**{name: data}, **kwargs))
 
+    @overload
     @classmethod
     def zeros(
         cls: Type[DataArrayClass[P, TDataArray]],
         shape: Union[Shape, Sizes],
         order: Order = "C",
         **kwargs: Any,
     ) -> TDataArray:
+        ...
+
+    @overload
+    @classmethod
+    def zeros(
+        cls: Type[DataClass[P]],
+        shape: Union[Shape, Sizes],
+        order: Order = "C",
+        **kwargs: Any,
+    ) -> xr.DataArray:
+        ...
+
+    @classmethod
+    def zeros(
+        cls: Any,
+        shape: Union[Shape, Sizes],
+        order: Order = "C",
+        **kwargs: Any,
+    ) -> Any:
         """Create a DataArray object filled with zeros.
 
         Args:
@@ -216,13 +262,33 @@ def zeros(
         data = np.zeros(shape, order=order)
         return asdataarray(cls(**{name: data}, **kwargs))
 
+    @overload
     @classmethod
     def ones(
         cls: Type[DataArrayClass[P, TDataArray]],
         shape: Union[Shape, Sizes],
         order: Order = "C",
         **kwargs: Any,
     ) -> TDataArray:
+        ...
+
+    @overload
+    @classmethod
+    def ones(
+        cls: Type[DataClass[P]],
+        shape: Union[Shape, Sizes],
+        order: Order = "C",
+        **kwargs: Any,
+    ) -> xr.DataArray:
+        ...
+
+    @classmethod
+    def ones(
+        cls: Any,
+        shape: Union[Shape, Sizes],
+        order: Order = "C",
+        **kwargs: Any,
+    ) -> Any:
         """Create a DataArray object filled with ones.
 
         Args:
@@ -244,6 +310,7 @@ def ones(
         data = np.ones(shape, order=order)
         return asdataarray(cls(**{name: data}, **kwargs))
 
+    @overload
     @classmethod
     def full(
         cls: Type[DataArrayClass[P, TDataArray]],
@@ -252,6 +319,27 @@ def full(
         order: Order = "C",
         **kwargs: Any,
     ) -> TDataArray:
+        ...
+
+    @overload
+    @classmethod
+    def full(
+        cls: Type[DataClass[P]],
+        shape: Union[Shape, Sizes],
+        fill_value: Any,
+        order: Order = "C",
+        **kwargs: Any,
+    ) -> xr.DataArray:
+        ...
+
+    @classmethod
+    def full(
+        cls: Any,
+        shape: Union[Shape, Sizes],
+        fill_value: Any,
+        order: Order = "C",
+        **kwargs: Any,
+    ) -> Any:
         """Create a DataArray object filled with given value.
 
         Args: