-
-
Notifications
You must be signed in to change notification settings - Fork 145
Description
Pandas version checks
-
I have checked that this issue has not already been reported.
-
I have confirmed this bug exists on the latest version of pandas.
-
I have confirmed this bug exists on the main branch of pandas.
Reproducible Example
from enum import Enum
import pandas as pd
class MyEnum(Enum):
TAYYAR = "tayyar"
HAYDAR = "haydar"
df = pd.DataFrame(data = [[12.2, 10], [8.8, 15], [22.1, 14]], columns=[MyEnum.TAYYAR, MyEnum.HAYDAR])
print(df[MyEnum.TAYYAR])
Issue Description
The example code above takes an Enum as column to a DataFrame and uses that enum in __get_item__
.
It works fine, which creates an inconsistency with the signature of __get_item__
method.
Here is the signature of __get_item__
. According to this signature, enums are not accepted as index.
def __getitem__(self, Union[str, bytes, date, datetime, timedelta, bool, int, float, complex, Timestamp, Timedelta], /) -> Series[Any]
When you use a static type checker e.g. mypy on the code example above, it fails.
mypy_enum.py:12: error: No overload variant of "__getitem__" of "DataFrame" matches argument type "MyEnum" [call-overload]
mypy_enum.py:12: note: Possible overload variants:
mypy_enum.py:12: note: def __getitem__(self, Union[str, bytes, date, datetime, timedelta, bool, int, float, complex, Timestamp, Timedelta], /) -> Series[Any]
mypy_enum.py:12: note: def __getitem__(self, slice, /) -> DataFrame
mypy_enum.py:12: note: def [ScalarT] __getitem__(self, Union[Tuple[Any, ...], Series[bool], DataFrame, List[str], List[ScalarT], Index, ndarray[Any, dtype[str_]], ndarray[Any, dtype[bool_]], Sequence[Tuple[Union[str, bytes, date, datetime, timedelta, bool, int, float, complex, Timestamp, Timedelta], ...]]], /) -> DataFrame
Found 1 error in 1 file (checked 1 source file)
Expected Behavior
__getitem__
's signature should be updated to support the column type itself. The column type is often not limited to primitive types. In can be custom objects as well.
Installed Versions
assert '_distutils' in core.__file__, core.__file__