122 lines
3.5 KiB
Python
122 lines
3.5 KiB
Python
""" define generic base classes for pandas objects """
|
|
from __future__ import annotations
|
|
|
|
from typing import TYPE_CHECKING, Type, cast
|
|
|
|
if TYPE_CHECKING:
|
|
from pandas import (
|
|
CategoricalIndex,
|
|
DataFrame,
|
|
DatetimeIndex,
|
|
Float64Index,
|
|
Int64Index,
|
|
IntervalIndex,
|
|
MultiIndex,
|
|
PeriodIndex,
|
|
RangeIndex,
|
|
Series,
|
|
TimedeltaIndex,
|
|
UInt64Index,
|
|
)
|
|
from pandas.core.generic import NDFrame
|
|
|
|
|
|
# define abstract base classes to enable isinstance type checking on our
|
|
# objects
|
|
def create_pandas_abc_type(name, attr, comp):
|
|
|
|
# https://github.com/python/mypy/issues/1006
|
|
# error: 'classmethod' used with a non-method
|
|
@classmethod # type: ignore[misc]
|
|
def _check(cls, inst) -> bool:
|
|
return getattr(inst, attr, "_typ") in comp
|
|
|
|
dct = {"__instancecheck__": _check, "__subclasscheck__": _check}
|
|
meta = type("ABCBase", (type,), dct)
|
|
return meta(name, (), dct)
|
|
|
|
|
|
ABCInt64Index = cast(
|
|
"Type[Int64Index]",
|
|
create_pandas_abc_type("ABCInt64Index", "_typ", ("int64index",)),
|
|
)
|
|
ABCUInt64Index = cast(
|
|
"Type[UInt64Index]",
|
|
create_pandas_abc_type("ABCUInt64Index", "_typ", ("uint64index",)),
|
|
)
|
|
ABCRangeIndex = cast(
|
|
"Type[RangeIndex]",
|
|
create_pandas_abc_type("ABCRangeIndex", "_typ", ("rangeindex",)),
|
|
)
|
|
ABCFloat64Index = cast(
|
|
"Type[Float64Index]",
|
|
create_pandas_abc_type("ABCFloat64Index", "_typ", ("float64index",)),
|
|
)
|
|
ABCMultiIndex = cast(
|
|
"Type[MultiIndex]",
|
|
create_pandas_abc_type("ABCMultiIndex", "_typ", ("multiindex",)),
|
|
)
|
|
ABCDatetimeIndex = cast(
|
|
"Type[DatetimeIndex]",
|
|
create_pandas_abc_type("ABCDatetimeIndex", "_typ", ("datetimeindex",)),
|
|
)
|
|
ABCTimedeltaIndex = cast(
|
|
"Type[TimedeltaIndex]",
|
|
create_pandas_abc_type("ABCTimedeltaIndex", "_typ", ("timedeltaindex",)),
|
|
)
|
|
ABCPeriodIndex = cast(
|
|
"Type[PeriodIndex]",
|
|
create_pandas_abc_type("ABCPeriodIndex", "_typ", ("periodindex",)),
|
|
)
|
|
ABCCategoricalIndex = cast(
|
|
"Type[CategoricalIndex]",
|
|
create_pandas_abc_type("ABCCategoricalIndex", "_typ", ("categoricalindex",)),
|
|
)
|
|
ABCIntervalIndex = cast(
|
|
"Type[IntervalIndex]",
|
|
create_pandas_abc_type("ABCIntervalIndex", "_typ", ("intervalindex",)),
|
|
)
|
|
ABCIndexClass = create_pandas_abc_type(
|
|
"ABCIndexClass",
|
|
"_typ",
|
|
{
|
|
"index",
|
|
"int64index",
|
|
"rangeindex",
|
|
"float64index",
|
|
"uint64index",
|
|
"multiindex",
|
|
"datetimeindex",
|
|
"timedeltaindex",
|
|
"periodindex",
|
|
"categoricalindex",
|
|
"intervalindex",
|
|
},
|
|
)
|
|
|
|
ABCNDFrame = cast(
|
|
"Type[NDFrame]",
|
|
create_pandas_abc_type("ABCNDFrame", "_typ", ("series", "dataframe")),
|
|
)
|
|
ABCSeries = cast(
|
|
"Type[Series]",
|
|
create_pandas_abc_type("ABCSeries", "_typ", ("series",)),
|
|
)
|
|
ABCDataFrame = cast(
|
|
"Type[DataFrame]", create_pandas_abc_type("ABCDataFrame", "_typ", ("dataframe",))
|
|
)
|
|
|
|
ABCCategorical = create_pandas_abc_type("ABCCategorical", "_typ", ("categorical"))
|
|
ABCDatetimeArray = create_pandas_abc_type("ABCDatetimeArray", "_typ", ("datetimearray"))
|
|
ABCTimedeltaArray = create_pandas_abc_type(
|
|
"ABCTimedeltaArray", "_typ", ("timedeltaarray")
|
|
)
|
|
ABCPeriodArray = create_pandas_abc_type("ABCPeriodArray", "_typ", ("periodarray",))
|
|
ABCExtensionArray = create_pandas_abc_type(
|
|
"ABCExtensionArray",
|
|
"_typ",
|
|
# Note: IntervalArray and SparseArray are included bc they have _typ="extension"
|
|
{"extension", "categorical", "periodarray", "datetimearray", "timedeltaarray"},
|
|
)
|
|
ABCPandasArray = create_pandas_abc_type("ABCPandasArray", "_typ", ("npy_extension",))
|