#121 Use overload in AsDataset

Author: astropenguin

xarray_dataclasses/dataset.py

@@ -49,32 +49,12 @@ def __init__(self, *args: P.args, **kwargs: P.kwargs) -> None:
     __dataoptions__: DataOptions[TDataset]
 
 
-# custom classproperty
-class classproperty:
-    """Class property only for AsDataset.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[DatasetClass[P, TDataset]],
-    ) -> Callable[P, TDataset]:
-        return self.__func__(cls)
-
-
 # runtime functions and classes
 @overload
 def asdataset(
     dataclass: DatasetClass[Any, TDataset],
     reference: Optional[DataType] = None,
-    dataoptions: Any = DEFAULT_OPTIONS,
+    dataoptions: DataOptions[Any] = DEFAULT_OPTIONS,
 ) -> TDataset:
     ...
 
@@ -90,8 +70,8 @@ def asdataset(
 
 def asdataset(
     dataclass: Any,
-    reference: Any = None,
-    dataoptions: Any = DEFAULT_OPTIONS,
+    reference: Optional[DataType] = None,
+    dataoptions: DataOptions[Any] = DEFAULT_OPTIONS,
 ) -> Any:
     """Create a Dataset object from a dataclass object.
 
@@ -135,36 +115,82 @@ def asdataset(
     return dataset
 
 
+# runtime classes
+class classproperty:
+    """Class property only for AsDataset.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
+
+    @overload
+    def __get__(
+        self,
+        obj: Any,
+        cls: Type[DatasetClass[P, TDataset]],
+    ) -> Callable[P, TDataset]:
+        ...
+
+    @overload
+    def __get__(
+        self,
+        obj: Any,
+        cls: Type[DataClass[P]],
+    ) -> Callable[P, xr.Dataset]:
+        ...
+
+    def __get__(self, obj: Any, cls: Any) -> Any:
+        return self.__func__(cls)
+
+
 class AsDataset:
     """Mix-in class that provides shorthand methods."""
 
-    __dataoptions__ = DEFAULT_OPTIONS
-
     @classproperty
-    def new(cls: Type[DatasetClass[P, TDataset]]) -> Callable[P, TDataset]:
+    def new(cls: Any) -> Any:
         """Create a Dataset 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"] = TDataset
 
         @wraps(init)
-        def new(
-            cls: Type[DatasetClass[P, TDataset]],
-            *args: P.args,
-            **kwargs: P.kwargs,
-        ) -> TDataset:
+        def new(cls: Any, *args: Any, **kwargs: Any) -> Any:
             return asdataset(cls(*args, **kwargs))
 
         return MethodType(new, cls)
 
+    @overload
     @classmethod
     def empty(
         cls: Type[DatasetClass[P, TDataset]],
         sizes: Sizes,
         order: Order = "C",
         **kwargs: Any,
     ) -> TDataset:
+        ...
+
+    @overload
+    @classmethod
+    def empty(
+        cls: Type[DataClass[P]],
+        sizes: Sizes,
+        order: Order = "C",
+        **kwargs: Any,
+    ) -> xr.Dataset:
+        ...
+
+    @classmethod
+    def empty(
+        cls: Any,
+        sizes: Sizes,
+        order: Order = "C",
+        **kwargs: Any,
+    ) -> Any:
         """Create a Dataset object without initializing data vars.
 
         Args:
@@ -186,13 +212,33 @@ def empty(
 
         return asdataset(cls(**data_vars, **kwargs))
 
+    @overload
     @classmethod
     def zeros(
         cls: Type[DatasetClass[P, TDataset]],
         sizes: Sizes,
         order: Order = "C",
         **kwargs: Any,
     ) -> TDataset:
+        ...
+
+    @overload
+    @classmethod
+    def zeros(
+        cls: Type[DataClass[P]],
+        sizes: Sizes,
+        order: Order = "C",
+        **kwargs: Any,
+    ) -> xr.Dataset:
+        ...
+
+    @classmethod
+    def zeros(
+        cls: Any,
+        sizes: Sizes,
+        order: Order = "C",
+        **kwargs: Any,
+    ) -> Any:
         """Create a Dataset object whose data vars are filled with zeros.
 
         Args:
@@ -214,13 +260,33 @@ def zeros(
 
         return asdataset(cls(**data_vars, **kwargs))
 
+    @overload
     @classmethod
     def ones(
         cls: Type[DatasetClass[P, TDataset]],
         sizes: Sizes,
         order: Order = "C",
         **kwargs: Any,
     ) -> TDataset:
+        ...
+
+    @overload
+    @classmethod
+    def ones(
+        cls: Type[DataClass[P]],
+        sizes: Sizes,
+        order: Order = "C",
+        **kwargs: Any,
+    ) -> xr.Dataset:
+        ...
+
+    @classmethod
+    def ones(
+        cls: Any,
+        sizes: Sizes,
+        order: Order = "C",
+        **kwargs: Any,
+    ) -> Any:
         """Create a Dataset object whose data vars are filled with ones.
 
         Args:
@@ -242,6 +308,7 @@ def ones(
 
         return asdataset(cls(**data_vars, **kwargs))
 
+    @overload
     @classmethod
     def full(
         cls: Type[DatasetClass[P, TDataset]],
@@ -250,6 +317,27 @@ def full(
         order: Order = "C",
         **kwargs: Any,
     ) -> TDataset:
+        ...
+
+    @overload
+    @classmethod
+    def full(
+        cls: Type[DataClass[P]],
+        sizes: Sizes,
+        fill_value: Any,
+        order: Order = "C",
+        **kwargs: Any,
+    ) -> xr.Dataset:
+        ...
+
+    @classmethod
+    def full(
+        cls: Any,
+        sizes: Sizes,
+        fill_value: Any,
+        order: Order = "C",
+        **kwargs: Any,
+    ) -> Any:
         """Create a Dataset object whose data vars are filled with given value.
 
         Args: