Repository: astropenguin/pandas-dataclasses

Issue:
Add generic version of mix-in class

Discussion:

Patch:
diff --git a/pandas_dataclasses/pandas/__init__.py b/pandas_dataclasses/pandas/__init__.py
index e183f05..5080d12 100644
--- a/pandas_dataclasses/pandas/__init__.py
+++ b/pandas_dataclasses/pandas/__init__.py
@@ -1,3 +1,4 @@
 # submodules
 from .frame import *
+from .generic import *
 from .series import *
diff --git a/pandas_dataclasses/pandas/frame.py b/pandas_dataclasses/pandas/frame.py
index 894c85f..a0ff9b2 100644
--- a/pandas_dataclasses/pandas/frame.py
+++ b/pandas_dataclasses/pandas/frame.py
@@ -23,12 +23,7 @@ TDataFrame = TypeVar("TDataFrame", bound=pd.DataFrame)
 
 # runtime classes
 class classproperty:
-    """Class property only for AsDataFrame.new().
-
-    As ``classmethod`` and ``property`` can be chained since Python 3.9,
-    this will be removed when the support for Python 3.7 and 3.8 ends.
-
-    """
+    """Create a DataFrame object from dataclass parameters."""
 
     def __init__(self, func: Any) -> None:
         self.__func__ = func
@@ -59,8 +54,9 @@ class AsDataFrame:
     @classproperty
     def new(cls) -> Any:
         """Create a DataFrame object from dataclass parameters."""
+        factory = getattr(cls, "__pandas_factory__", pd.DataFrame)
         init = copy(cls.__init__)
-        init.__annotations__["return"] = pd.DataFrame
+        init.__annotations__["return"] = factory
 
         @wraps(init)
         def new(cls: Any, *args: Any, **kwargs: Any) -> Any:
@@ -71,17 +67,17 @@ class AsDataFrame:
 
 # runtime functions
 @overload
-def asdataframe(obj: PandasClass[P, TDataFrame], *, factory: None = None) -> TDataFrame:
+def asdataframe(obj: Any, *, factory: Type[TDataFrame]) -> TDataFrame:
     ...
 
 
 @overload
-def asdataframe(obj: DataClass[P], *, factory: None = None) -> pd.DataFrame:
+def asdataframe(obj: PandasClass[P, TDataFrame], *, factory: None = None) -> TDataFrame:
     ...
 
 
 @overload
-def asdataframe(obj: Any, *, factory: Type[TDataFrame] = pd.DataFrame) -> TDataFrame:
+def asdataframe(obj: DataClass[P], *, factory: None = None) -> pd.DataFrame:
     ...
 
 
diff --git a/pandas_dataclasses/pandas/generic.py b/pandas_dataclasses/pandas/generic.py
new file mode 100644
index 0000000..f6b6677
--- /dev/null
+++ b/pandas_dataclasses/pandas/generic.py
@@ -0,0 +1,82 @@
+__all__ = ["As"]
+
+
+# standard library
+from functools import wraps
+from types import MethodType
+from typing import Any, Callable, Generic, Type, TypeVar
+
+
+# dependencies
+import pandas as pd
+from morecopy import copy
+from typing_extensions import get_args, get_origin
+
+
+# submodules
+from .frame import asdataframe
+from .series import asseries
+from ..typing import AnyPandas, P, PandasClass
+
+
+# type hints
+TPandas = TypeVar("TPandas", bound=AnyPandas)
+
+
+# runtime classes
+class classproperty:
+    """Create a pandas object from dataclass parameters."""
+
+    def __init__(self, func: Any) -> None:
+        self.__func__ = func
+
+    def __get__(
+        self,
+        obj: Any,
+        cls: Type[PandasClass[P, TPandas]],
+    ) -> Callable[P, TPandas]:
+        return self.__func__(cls)
+
+
+class As(Generic[TPandas]):
+    """Mix-in class that provides shorthand methods."""
+
+    __pandas_factory__: Type[TPandas]
+
+    def __init_subclass__(cls, **kwargs: Any) -> None:
+        """Add a pandas factory to an inheriting class."""
+        cls.__pandas_factory__ = get_factory(cls)
+        return super().__init_subclass__(**kwargs)
+
+    @classproperty
+    def new(cls) -> Any:
+        """Create a pandas object from dataclass parameters."""
+        factory = cls.__pandas_factory__
+        init = copy(cls.__init__)
+        init.__annotations__["return"] = factory
+
+        if issubclass(factory, pd.DataFrame):
+            aspandas = asdataframe
+        else:
+            aspandas = asseries
+
+        @wraps(init)
+        def new(cls: Any, *args: Any, **kwargs: Any) -> Any:
+            return aspandas(cls(*args, **kwargs))
+
+        return MethodType(new, cls)
+
+
+# runtime functions
+def get_factory(cls: Any) -> Type[AnyPandas]:
+    """Extract a pandas factory from an As-inherited class."""
+    for base in cls.__orig_bases__:
+        if get_origin(base) is not As:
+            continue
+
+        factory = get_args(base)[0]
+
+        if issubclass(factory, (pd.DataFrame, pd.Series)):
+            return factory
+
+    raise TypeError("Could not find any pandas factory.")
diff --git a/pandas_dataclasses/pandas/series.py b/pandas_dataclasses/pandas/series.py
index d5236ed..1875376 100644
--- a/pandas_dataclasses/pandas/series.py
+++ b/pandas_dataclasses/pandas/series.py
@@ -23,12 +23,7 @@ TSeries = TypeVar("TSeries", bound=pd.Series)
 
 # runtime classes
 class classproperty:
-    """Class property only for AsSeries.new().
-
-    As ``classmethod`` and ``property`` can be chained since Python 3.9,
-    this will be removed when the support for Python 3.7 and 3.8 ends.
-
-    """
+    """Create a Series object from dataclass parameters."""
 
     def __init__(self, func: Any) -> None:
         self.__func__ = func
@@ -59,8 +54,9 @@ class AsSeries:
     @classproperty
     def new(cls) -> Any:
         """Create a Series object from dataclass parameters."""
+        factory = getattr(cls, "__pandas_factory__", pd.Series)
         init = copy(cls.__init__)
-        init.__annotations__["return"] = pd.Series
+        init.__annotations__["return"] = factory
 
         @wraps(init)
         def new(cls: Any, *args: Any, **kwargs: Any) -> Any:
@@ -71,17 +67,17 @@ class AsSeries:
 
 # runtime functions
 @overload
-def asseries(obj: PandasClass[P, TSeries], *, factory: None = None) -> TSeries:
+def asseries(obj: Any, *, factory: Type[TSeries]) -> TSeries:
     ...
 
 
 @overload
-def asseries(obj: DataClass[P], *, factory: None = None) -> pd.Series:
+def asseries(obj: PandasClass[P, TSeries], *, factory: None = None) -> TSeries:
     ...
 
 
 @overload
-def asseries(obj: Any, *, factory: Type[TSeries] = pd.Series) -> TSeries:
+def asseries(obj: DataClass[P], *, factory: None = None) -> pd.Series:
     ...
 
 


Test patch:
diff --git a/tests/test_frame_generic.py b/tests/test_frame_generic.py
new file mode 100644
index 0000000..fec6891
--- /dev/null
+++ b/tests/test_frame_generic.py
@@ -0,0 +1,74 @@
+# standard library
+from dataclasses import dataclass
+
+
+# dependencies
+import numpy as np
+import pandas as pd
+from pandas_dataclasses.pandas.generic import As
+from pandas_dataclasses.typing import Attr, Data, Index
+from typing_extensions import Annotated as Ann
+
+
+# test datasets
+class CustomFrame(pd.DataFrame):
+    """Custom pandas DataFrame."""
+
+    __slots__ = ()
+
+
+@dataclass
+class Weather(As[CustomFrame]):
+    """Weather information at a location."""
+
+    year: Ann[Index[int], "Year"]
+    month: Ann[Index[int], "Month"]
+    temp: Ann[Data[float], "Average temperature (deg C)"]
+    humid: Ann[Data[float], "Average humidity (%)"]
+    loc: Ann[Attr[str], "Location"] = "Tokyo"
+    lon: Ann[Attr[float], "Longitude (deg)"] = 139.69167
+    lat: Ann[Attr[float], "Latitude (deg)"] = 35.68944
+
+
+year = np.array([2020, 2020, 2021, 2021, 2022])
+month = np.array([1, 7, 1, 7, 1])
+temp = np.array([7.1, 24.3, 5.4, 25.9, 4.9])
+humid = np.array([65, 89, 57, 83, 52])
+weather = Weather.new(year, month, temp, humid)
+
+
+# test functions
+def test_attrs() -> None:
+    assert weather.attrs == {
+        "Location": "Tokyo",
+        "Longitude (deg)": 139.69167,
+        "Latitude (deg)": 35.68944,
+    }
+
+
+def test_data() -> None:
+    data_0 = weather.iloc[:, 0]
+    data_1 = weather.iloc[:, 1]
+
+    assert (data_0 == temp).all()
+    assert (data_1 == humid).all()
+    assert data_0.dtype == np.dtype("float64")
+    assert data_1.dtype == np.dtype("float64")
+    assert data_0.name == "Average temperature (deg C)"
+    assert data_1.name == "Average humidity (%)"
+
+
+def test_index() -> None:
+    index_0 = weather.index.get_level_values(0)
+    index_1 = weather.index.get_level_values(1)
+
+    assert (index_0 == year).all()
+    assert (index_1 == month).all()
+    assert index_0.dtype == np.dtype("int64")
+    assert index_1.dtype == np.dtype("int64")
+    assert index_0.name == "Year"
+    assert index_1.name == "Month"
+
+
+def test_instance() -> None:
+    assert isinstance(weather, CustomFrame)
diff --git a/tests/test_series_generic.py b/tests/test_series_generic.py
new file mode 100644
index 0000000..f0d5be0
--- /dev/null
+++ b/tests/test_series_generic.py
@@ -0,0 +1,62 @@
+# standard library
+from dataclasses import dataclass
+
+
+# dependencies
+import numpy as np
+import pandas as pd
+from pandas_dataclasses.pandas.generic import As
+from pandas_dataclasses.typing import Attr, Data, Index, Name
+from typing_extensions import Annotated as Ann
+from typing_extensions import Literal as L
+
+
+# test datasets
+class CustomSeries(pd.Series):
+    """Custom pandas Series."""
+
+    __slots__ = ()
+
+
+@dataclass
+class Temperature(As[CustomSeries]):
+    """Temperature information at a location."""
+
+    time: Ann[Index[L["M8[ns]"]], "Time in UTC"]
+    data: Data[float]
+    loc: Ann[Attr[str], "Location"] = "Tokyo"
+    lon: Ann[Attr[float], "Longitude (deg)"] = 139.69167
+    lat: Ann[Attr[float], "Latitude (deg)"] = 35.68944
+    name: Name[str] = "Temperature (deg C)"
+
+
+time = pd.date_range("2020-01", "2020-06", freq="MS")
+data = np.array([7.1, 8.3, 10.7, 12.8, 19.5, 23.2])
+temperature = Temperature.new(time, data)
+
+
+# test functions
+def test_attrs() -> None:
+    assert temperature.attrs == {
+        "Location": "Tokyo",
+        "Longitude (deg)": 139.69167,
+        "Latitude (deg)": 35.68944,
+    }
+
+
+def test_data() -> None:
+    assert temperature.to_numpy() is data
+    assert temperature.name == "Temperature (deg C)"
+    assert temperature.dtype == np.float64
+
+
+def test_index() -> None:
+    index = temperature.index
+
+    assert index.to_numpy() is time.to_numpy()
+    assert index.name == "Time in UTC"
+    assert index.dtype == np.dtype("M8[ns]")
+
+
+def test_instance() -> None:
+    assert isinstance(temperature, CustomSeries)