From b9596acb19378b7e378f0d522bc4ec77579a6d04 Mon Sep 17 00:00:00 2001 From: Anil Tuncel Date: Tue, 4 Apr 2023 23:57:01 +0200 Subject: [PATCH] update signature of Series.pop #627 --- pandas-stubs/core/series.pyi | 2 +- tests/test_series.py | 16 ++++++++++++++++ 2 files changed, 17 insertions(+), 1 deletion(-) diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index 8ffbd5df..3ff1f07e 100644 --- a/pandas-stubs/core/series.pyi +++ b/pandas-stubs/core/series.pyi @@ -1021,7 +1021,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]): def droplevel( self, level: Level | list[Level], axis: AxisIndex = ... ) -> DataFrame: ... - def pop(self, item: _str) -> Series[S1]: ... + def pop(self, item: Hashable) -> Any: ... def squeeze(self, axis: AxisIndex | None = ...) -> Scalar: ... def __abs__(self) -> Series[S1]: ... def add_prefix(self, prefix: _str) -> Series[S1]: ... diff --git a/tests/test_series.py b/tests/test_series.py index 68190038..2e8b665f 100644 --- a/tests/test_series.py +++ b/tests/test_series.py @@ -2,6 +2,7 @@ import datetime from decimal import Decimal +from enum import Enum from pathlib import Path import re from typing import ( @@ -226,6 +227,21 @@ def test_types_dropna() -> None: assert assert_type(s.dropna(axis=0, inplace=True), None) is None +def test_pop() -> None: + # Testing pop support for hashable types + # Due to the bug in https://github.com/pandas-dev/pandas-stubs/issues/627 + class MyEnum(Enum): + FIRST = "tayyar" + SECOND = "haydar" + + df = pd.DataFrame( + data=[[12.2, 10], [8.8, 15]], columns=[MyEnum.FIRST, MyEnum.SECOND] + ) + series = df.loc[0] + res = series.pop(MyEnum.FIRST) + check(assert_type(res, Any), float) + + def test_types_fillna() -> None: s = pd.Series([1, np.nan, np.nan, 3]) check(assert_type(s.fillna(0), pd.Series), pd.Series)