Repository: astropenguin/pandas-dataclasses

Issue:
Extend pandas factory
Update type hints for pandas factory (`__pandas_factory__`) to accept a function whose return type is a subclass of `pandas.DataFrame` or `pandas.Series`.

Comments:

Patch:
diff --git a/pandas_dataclasses/core/asdata.py b/pandas_dataclasses/core/asdata.py
index 6e0fb59..29265a5 100644
--- a/pandas_dataclasses/core/asdata.py
+++ b/pandas_dataclasses/core/asdata.py
@@ -2,7 +2,7 @@ __all__ = ["asdataframe", "asseries"]
 
 
 # standard library
-from typing import Any, Dict, Hashable, List, Optional, Type, overload
+from typing import Any, Callable, Dict, Hashable, List, Optional, overload
 
 
 # dependencies
@@ -17,12 +17,12 @@ from .typing import P, DataClass, PandasClass, TDataFrame, TSeries
 
 # runtime functions
 @overload
-def asdataframe(obj: Any, *, factory: Type[TDataFrame]) -> TDataFrame:
+def asdataframe(obj: PandasClass[P, TDataFrame], *, factory: None = None) -> TDataFrame:
     ...
 
 
 @overload
-def asdataframe(obj: PandasClass[P, TDataFrame], *, factory: None = None) -> TDataFrame:
+def asdataframe(obj: DataClass[P], *, factory: Callable[..., TDataFrame]) -> TDataFrame:
     ...
 
 
@@ -38,9 +38,6 @@ def asdataframe(obj: Any, *, factory: Any = None) -> Any:
     if factory is None:
         factory = spec.factory or pd.DataFrame
 
-    if not issubclass(factory, pd.DataFrame):
-        raise TypeError("Factory must be a subclass of DataFrame.")
-
     dataframe = factory(
         data=get_data(spec),
         index=get_index(spec),
@@ -52,12 +49,12 @@ def asdataframe(obj: Any, *, factory: Any = None) -> Any:
 
 
 @overload
-def asseries(obj: Any, *, factory: Type[TSeries]) -> TSeries:
+def asseries(obj: PandasClass[P, TSeries], *, factory: None = None) -> TSeries:
     ...
 
 
 @overload
-def asseries(obj: PandasClass[P, TSeries], *, factory: None = None) -> TSeries:
+def asseries(obj: DataClass[P], *, factory: Callable[..., TSeries]) -> TSeries:
     ...
 
 
@@ -73,9 +70,6 @@ def asseries(obj: Any, *, factory: Any = None) -> Any:
     if factory is None:
         factory = spec.factory or pd.Series
 
-    if not issubclass(factory, pd.Series):
-        raise TypeError("Factory must be a subclass of Series.")
-
     data = get_data(spec)
     index = get_index(spec)
 
diff --git a/pandas_dataclasses/core/mixins.py b/pandas_dataclasses/core/mixins.py
index f900957..0e21fe2 100644
--- a/pandas_dataclasses/core/mixins.py
+++ b/pandas_dataclasses/core/mixins.py
@@ -5,7 +5,7 @@ __all__ = ["As", "AsDataFrame", "AsSeries"]
 from copy import copy
 from functools import wraps as wraps_
 from types import FunctionType, MethodType
-from typing import Any, Callable, ForwardRef, Generic, Type, cast
+from typing import Any, Callable, ForwardRef, Generic, Type
 
 
 # dependencies
@@ -15,13 +15,13 @@ from typing_extensions import get_args, get_origin
 
 # submodules
 from .asdata import asdataframe, asseries
-from .typing import P, T, PandasClass, TPandas
+from .typing import P, T, Pandas, PandasClass, TPandas
 
 
 class classproperty:
-    """Class property dedicated to ``As.new``."""
+    """Class property decorator dedicated to ``As.new``."""
 
-    def __init__(self, func: Any) -> None:
+    def __init__(self, func: Callable[..., Any]) -> None:
         self.__doc__ = func.__doc__
         self.func = func
 
@@ -33,20 +33,82 @@ class classproperty:
         return self.func(cls)  # type: ignore
 
 
-def wraps(func: Any, return_: Any) -> Callable[[T], T]:
-    """Function decorator dedicated to ``As.new``."""
+class As(Generic[TPandas]):
+    """Mix-in class for runtime pandas data creator."""
+
+    __pandas_factory__: Callable[..., TPandas]
+    """Factory for pandas data creation."""
+
+    def __init_subclass__(cls, **kwargs: Any) -> None:
+        """Add a pandas factory to an inheriting class."""
+        factory = kwargs.pop("factory", None)
+        cls.__pandas_factory__ = factory or get_factory(cls)
+        super().__init_subclass__(**kwargs)
+
+    @classproperty
+    def new(cls) -> MethodType:
+        """Runtime pandas data creator as a classmethod."""
+        return MethodType(get_new(cls), cls)
+
+
+AsDataFrame = As[pd.DataFrame]
+"""Alias of ``As[pandas.DataFrame]``."""
+
+
+AsSeries = As["pd.Series[Any]"]
+"""Alias of ``As[pandas.Series[Any]]``."""
+
+
+def get_factory(cls: Any) -> Callable[..., Any]:
+    """Extract a pandas factory from a class."""
+    for base in getattr(cls, "__orig_bases__", ()):
+        if get_origin(base) is not As:
+            continue
+
+        factory = get_args(base)[0]
+
+        # special handling for AsSeries
+        if factory == ForwardRef("pd.Series[Any]"):
+            return pd.Series
+
+        return factory  # type: ignore
+
+    raise TypeError("Could not find any factory.")
+
+
+def get_new(cls: Any) -> Callable[..., Pandas]:
+    """Create a runtime new function from a class."""
+    factory = cls.__pandas_factory__
+    origin = get_origin(factory) or factory
+
+    if issubclass(origin, pd.DataFrame):
+        converter: Any = asdataframe
+    elif issubclass(origin, pd.Series):
+        converter = asseries
+    else:
+        raise TypeError("Could not choose a converter.")
+
+    @wraps(cls.__init__, "new", factory)
+    def new(cls: Any, *args: Any, **kwargs: Any) -> Any:
+        return converter(cls(*args, **kwargs))
+
+    return new
+
+
+def wraps(func: Any, name: str, return_: Any) -> Callable[[T], T]:
+    """functools.wraps with modifiable name and return type."""
     if not isinstance(func, FunctionType):
         return wraps_(func)
 
     copied = type(func)(
         func.__code__,
         func.__globals__,
-        "new",
+        name,
         func.__defaults__,
         func.__closure__,
     )
 
-    for name in (
+    for attr in (
         "__annotations__",
         "__dict__",
         "__doc__",
@@ -55,55 +117,7 @@ def wraps(func: Any, return_: Any) -> Callable[[T], T]:
         "__name__",
         "__qualname__",
     ):
-        setattr(copied, name, copy(getattr(func, name)))
+        setattr(copied, attr, copy(getattr(func, attr)))
 
     copied.__annotations__["return"] = return_
     return wraps_(copied)
-
-
-class As(Generic[TPandas]):
-    """Mix-in class that provides shorthand methods."""
-
-    __pandas_factory__: Type[TPandas]
-    """Factory for pandas data creation."""
-
-    def __init_subclass__(cls, **kwargs: Any) -> None:
-        """Add a pandas factory to an inheriting class."""
-        super().__init_subclass__(**kwargs)
-
-        for base in cls.__orig_bases__:  # type: ignore
-            if get_origin(base) is not As:
-                continue
-
-            factory = get_args(base)[0]
-
-            if factory == ForwardRef("pd.Series[Any]"):
-                cls.__pandas_factory__ = cast(Any, pd.Series)
-            else:
-                cls.__pandas_factory__ = factory
-
-    @classproperty
-    def new(cls) -> Any:
-        """Create a pandas object from dataclass parameters."""
-        factory = cls.__pandas_factory__
-
-        if issubclass(factory, pd.DataFrame):
-            aspandas: Any = asdataframe
-        elif issubclass(factory, pd.Series):
-            aspandas = asseries
-        else:
-            raise TypeError("Not a valid pandas factory.")
-
-        @wraps(cls.__init__, factory)  # type: ignore
-        def new(cls: Any, *args: Any, **kwargs: Any) -> Any:
-            return aspandas(cls(*args, **kwargs))
-
-        return MethodType(new, cls)
-
-
-AsDataFrame = As[pd.DataFrame]
-"""Alias of ``As[pandas.DataFrame]``."""
-
-
-AsSeries = As["pd.Series[Any]"]
-"""Alias of ``As[pandas.Series[Any]]``."""
diff --git a/pandas_dataclasses/core/specs.py b/pandas_dataclasses/core/specs.py
index f33b5f5..009a1a3 100644
--- a/pandas_dataclasses/core/specs.py
+++ b/pandas_dataclasses/core/specs.py
@@ -5,7 +5,7 @@ __all__ = ["Spec"]
 from dataclasses import dataclass, replace
 from dataclasses import Field as Field_, fields as fields_
 from functools import lru_cache
-from typing import Any, Hashable, List, Optional, Type, Union
+from typing import Any, Callable, Hashable, List, Optional, Type
 
 
 # dependencies
@@ -13,8 +13,7 @@ from typing_extensions import Literal, get_type_hints
 
 
 # submodules
-import pandas as pd
-from .typing import P, DataClass, Role, get_dtype, get_name, get_role
+from .typing import P, DataClass, Pandas, Role, get_dtype, get_name, get_role
 
 
 # runtime classes
@@ -84,7 +83,7 @@ class Spec:
     fields: Fields
     """List of field specifications."""
 
-    factory: Optional[Type[Union[pd.DataFrame, "pd.Series[Any]"]]] = None
+    factory: Optional[Callable[..., Pandas]] = None
     """Factory for pandas data creation."""
 
     @classmethod
diff --git a/pandas_dataclasses/core/typing.py b/pandas_dataclasses/core/typing.py
index 8d3a1e2..dc0704a 100644
--- a/pandas_dataclasses/core/typing.py
+++ b/pandas_dataclasses/core/typing.py
@@ -7,13 +7,13 @@ from enum import Enum, auto
 from itertools import chain
 from typing import (
     Any,
+    Callable,
     Collection,
     Dict,
     Hashable,
     Iterable,
     Optional,
     Tuple,
-    Type,
     TypeVar,
     Union,
 )
@@ -34,9 +34,10 @@ from typing_extensions import (
 
 
 # type hints (private)
+Pandas = Union[pd.DataFrame, "pd.Series[Any]"]
 P = ParamSpec("P")
 T = TypeVar("T")
-TPandas = TypeVar("TPandas", bound=Union[pd.DataFrame, "pd.Series[Any]"])
+TPandas = TypeVar("TPandas", bound=Pandas)
 TDataFrame = TypeVar("TDataFrame", bound=pd.DataFrame)
 TSeries = TypeVar("TSeries", bound="pd.Series[Any]")
 
@@ -54,7 +55,7 @@ class PandasClass(Protocol[P, TPandas]):
     """Type hint for dataclass objects with a pandas factory."""
 
     __dataclass_fields__: Dict[str, "Field[Any]"]
-    __pandas_factory__: Type[TPandas]
+    __pandas_factory__: Callable[..., TPandas]
 
     def __init__(self, *args: P.args, **kwargs: P.kwargs) -> None:
         ...


Test patch:
diff --git a/tests/test_mixins.py b/tests/test_mixins.py
index 13cffc5..3abcaf2 100644
--- a/tests/test_mixins.py
+++ b/tests/test_mixins.py
@@ -24,12 +24,12 @@ class DataFrameWeather(Weather, AsDataFrame):
 
 
 @dataclass
-class SeriesWeather(Weather, AsSeries):
+class CustomDataFrameWeather(Weather, As[CustomDataFrame]):
     pass
 
 
 @dataclass
-class CustomDataFrameWeather(Weather, As[CustomDataFrame]):
+class SeriesWeather(Weather, AsSeries):
     pass
 
 
@@ -38,6 +38,11 @@ class CustomSeriesWeather(Weather, As[CustomSeries]):
     pass
 
 
+@dataclass
+class FloatSeriesWeather(Weather, As["pd.Series[float]"], factory=pd.Series):
+    pass
+
+
 # test functions
 def test_dataframe_weather() -> None:
     df_weather = DataFrameWeather.new(
@@ -93,3 +98,17 @@ def test_custom_series_weather() -> None:
 
     assert isinstance(ser_weather, CustomSeries)
     assert_series_equal(ser_weather, ser_weather_true, check_series_type=False)
+
+
+def test_float_series_weather() -> None:
+    ser_weather = FloatSeriesWeather.new(
+        year=weather.year,
+        month=weather.month,
+        temp_avg=weather.temp_avg,
+        temp_max=weather.temp_max,
+        wind_avg=weather.wind_avg,
+        wind_max=weather.wind_max,
+    )
+
+    assert isinstance(ser_weather, pd.Series)
+    assert_series_equal(ser_weather, ser_weather_true, check_series_type=False)