pandas-dataclasses-generic-as.py

# standard library
from dataclasses import dataclass
from typing import (
    Any,
    Callable,
    Generic,
    ParamSpec,
    Protocol,
    Type,
    TypeVar,
    get_args,
)


# dependencies
import pandas as pd


# type hints
P = ParamSpec("P")
AnyPandas = pd.DataFrame | pd.Series
TPandas = TypeVar("TPandas", bound=AnyPandas)


class PandasClass(Protocol[P, TPandas]):
    __pandas_factory__: Type[TPandas]

    def __init__(self, *args: P.args, **kwargs: P.kwargs) -> None:
        ...


# runtime classes
class classproperty:
    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]):
    __pandas_factory__: Type[TPandas]

    def __init_subclass__(cls) -> None:
        cls.__pandas_factory__ = get_factory(cls)
        return super().__init_subclass__()

    @classproperty
    def new(cls: Any) -> Any:
        return cls


# runtime functions
def get_factory(cls: Any) -> Type[AnyPandas]:
    return get_args(cls.__orig_bases__[0])[0]


if __name__ == "__main__":

    @dataclass
    class A(As[pd.Series]):
        a: int
        b: int

    @dataclass
    class B(As[pd.DataFrame]):
        a: int
        b: int

    ser = A.new(0, 1)
    df = B.new(0, 1)