#237 Add api module for dataclass converters

Author: astropenguin

xarray_dataclasses/__init__.py

@@ -1,5 +1,6 @@
 __all__ = [
     # submodules
+    "api",
     "typing",
     # aliases
     "Attr",
@@ -11,13 +12,18 @@
     "Factory",
     "Name",
     "Tag",
+    "asdataarray",
+    "asdataset",
+    "asxarray",
 ]
 __version__ = "2.0.0"
 
 
 # submodules
+from . import api
 from . import typing
 
 
 # aliases
+from .api import *
 from .typing import *

xarray_dataclasses/api.py

@@ -0,0 +1,97 @@
+__all__ = ["asdataarray", "asdataset", "asxarray"]
+
+
+# standard library
+from typing import Any, Callable, overload
+
+
+# dependencies
+from xarray import DataArray, Dataset
+from .typing import DataClass, DataClassOf, PAny, TDataArray, TDataset, TXarray
+
+
+@overload
+def asdataarray(
+    obj: DataClassOf[PAny, TDataArray],
+    /,
+    *,
+    factory: None = None,
+) -> TDataArray: ...
+
+
+@overload
+def asdataarray(
+    obj: DataClass[PAny],
+    /,
+    *,
+    factory: Callable[..., TDataArray],
+) -> TDataArray: ...
+
+
+@overload
+def asdataarray(
+    obj: DataClass[PAny],
+    /,
+    *,
+    factory: None = None,
+) -> DataArray: ...
+
+
+def asdataarray(obj: Any, /, *, factory: Any = None) -> Any:
+    """Create a DataArray object from a dataclass object."""
+    ...
+
+
+@overload
+def asdataset(
+    obj: DataClassOf[PAny, TDataset],
+    /,
+    *,
+    factory: None = None,
+) -> TDataset: ...
+
+
+@overload
+def asdataset(
+    obj: DataClass[PAny],
+    /,
+    *,
+    factory: Callable[..., TDataset],
+) -> TDataset: ...
+
+
+@overload
+def asdataset(
+    obj: DataClass[PAny],
+    /,
+    *,
+    factory: None = None,
+) -> Dataset: ...
+
+
+def asdataset(obj: Any, /, *, factory: Any = None) -> Any:
+    """Create a Dataset object from a dataclass object."""
+    ...
+
+
+@overload
+def asxarray(
+    obj: DataClassOf[PAny, TXarray],
+    /,
+    *,
+    factory: None = None,
+) -> TXarray: ...
+
+
+@overload
+def asxarray(
+    obj: DataClass[PAny],
+    /,
+    *,
+    factory: Callable[..., TXarray],
+) -> TXarray: ...
+
+
+def asxarray(obj: Any, /, *, factory: Any = None) -> Any:
+    """Create a DataArray/set object from a dataclass object."""
+    ...

xarray_dataclasses/typing.py

@@ -13,8 +13,18 @@
 
 # standard library
 from collections.abc import Collection as Collection_, Hashable
+from dataclasses import Field
 from enum import auto
-from typing import Annotated, Callable, Protocol, TypeVar, Union
+from typing import (
+    Annotated,
+    Any,
+    Callable,
+    ClassVar,
+    ParamSpec,
+    Protocol,
+    TypeVar,
+    Union,
+)
 
 
 # dependencies
@@ -23,11 +33,14 @@
 
 
 # type hints
+PAny = ParamSpec("PAny")
 TAny = TypeVar("TAny")
+TDataArray = TypeVar("TDataArray", bound=DataArray)
+TDataset = TypeVar("TDataset", bound=Dataset)
 TDims = TypeVar("TDims", covariant=True)
 TDtype = TypeVar("TDtype", covariant=True)
 THashable = TypeVar("THashable", bound=Hashable)
-TXarray = TypeVar("TXarray", bound="Xarray")
+TXarray = TypeVar("TXarray", covariant=True, bound="Xarray")
 Xarray = Union[DataArray, Dataset]
 
 
@@ -37,6 +50,23 @@ class Collection(Collection_[TDtype], Protocol[TDims, TDtype]):
     pass
 
 
+class DataClass(Protocol[PAny]):
+    """Protocol for a dataclass object."""
+
+    __dataclass_fields__: ClassVar[dict[str, Field[Any]]]
+
+    def __init__(self, *args: PAny.args, **kwargs: PAny.kwargs) -> None: ...
+
+
+class DataClassOf(Protocol[PAny, TXarray]):
+    """Protocol for a dataclass object with an xarray factory."""
+
+    _xarray_factory: Callable[..., TXarray]
+    __dataclass_fields__: ClassVar[dict[str, Field[Any]]]
+
+    def __init__(self, *args: PAny.args, **kwargs: PAny.kwargs) -> None: ...
+
+
 # constants
 class Tag(TagBase):
     """Collection of xarray-related tags for annotating type hints."""