#100 Make Series generic

Author: astropenguin

pandas_dataclasses/core/asdata.py

@@ -62,7 +62,7 @@ def asseries(obj: PandasClass[P, TSeries], *, factory: None = None) -> TSeries:
 
 
 @overload
-def asseries(obj: DataClass[P], *, factory: None = None) -> pd.Series:
+def asseries(obj: DataClass[P], *, factory: None = None) -> "pd.Series[Any]":
     ...
 
 

pandas_dataclasses/core/mixins.py

@@ -5,7 +5,7 @@
 from copy import copy
 from functools import wraps as wraps_
 from types import FunctionType, MethodType
-from typing import Any, Callable, Generic, Type
+from typing import Any, Callable, ForwardRef, Generic, Type, cast
 
 
 # dependencies
@@ -72,8 +72,15 @@ def __init_subclass__(cls, **kwargs: Any) -> None:
         super().__init_subclass__(**kwargs)
 
         for base in cls.__orig_bases__:  # type: ignore
-            if get_origin(base) is As:
-                cls.__pandas_factory__ = get_args(base)[0]
+            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:
@@ -98,5 +105,5 @@ def new(cls: Any, *args: Any, **kwargs: Any) -> Any:
 """Alias of ``As[pandas.DataFrame]``."""
 
 
-AsSeries = As[pd.Series]
-"""Alias of ``As[pandas.Series]``."""
+AsSeries = As["pd.Series[Any]"]
+"""Alias of ``As[pandas.Series[Any]]``."""

pandas_dataclasses/core/specs.py

@@ -84,7 +84,7 @@ class Spec:
     fields: Fields
     """List of field specifications."""
 
-    factory: Optional[Type[Union[pd.DataFrame, pd.Series]]] = None
+    factory: Optional[Type[Union[pd.DataFrame, "pd.Series[Any]"]]] = None
     """Factory for pandas data creation."""
 
     @classmethod

pandas_dataclasses/core/typing.py

@@ -36,9 +36,9 @@
 # type hints (private)
 P = ParamSpec("P")
 T = TypeVar("T")
-TPandas = TypeVar("TPandas", bound=Union[pd.DataFrame, pd.Series])
+TPandas = TypeVar("TPandas", bound=Union[pd.DataFrame, "pd.Series[Any]"])
 TDataFrame = TypeVar("TDataFrame", bound=pd.DataFrame)
-TSeries = TypeVar("TSeries", bound=pd.Series)
+TSeries = TypeVar("TSeries", bound="pd.Series[Any]")
 
 
 class DataClass(Protocol[P]):