Skip to content

Commit

Permalink
#100 Make Series generic
Browse files Browse the repository at this point in the history
  • Loading branch information
astropenguin committed Oct 10, 2022
1 parent 1cedf5f commit 958eb35
Show file tree
Hide file tree
Showing 4 changed files with 16 additions and 9 deletions.
2 changes: 1 addition & 1 deletion pandas_dataclasses/core/asdata.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]":
...


Expand Down
17 changes: 12 additions & 5 deletions pandas_dataclasses/core/mixins.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


# dependencies
Expand Down Expand Up @@ -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]"):
factory = pd.Series

cls.__pandas_factory__ = factory # type: ignore

@classproperty
def new(cls) -> Any:
Expand All @@ -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]]``."""
2 changes: 1 addition & 1 deletion pandas_dataclasses/core/specs.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
4 changes: 2 additions & 2 deletions pandas_dataclasses/core/typing.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]):
Expand Down

0 comments on commit 958eb35

Please sign in to comment.