Add get method to Series

Author: skatsuta

pandas-stubs/core/series.pyi

@@ -18,6 +18,7 @@ from typing import (
     ClassVar,
     Generic,
     Literal,
+    TypeVar,
     overload,
 )
 
@@ -160,6 +161,8 @@ from pandas.plotting import PlotAccessor
 _bool = bool
 _str = str
 
+_T = TypeVar("_T")
+
 class _iLocIndexerSeries(_iLocIndexer, Generic[S1]):
     # get item
     @overload
@@ -381,6 +384,10 @@ class Series(IndexOpsMixin[S1], NDFrame):
     @overload
     def __getitem__(self, idx: Scalar) -> S1: ...
     def __setitem__(self, key, value) -> None: ...
+    @overload
+    def get(self, key: Hashable, default: None = ...) -> S1 | None: ...
+    @overload
+    def get(self, key: Hashable, default: S1 | _T = ...) -> S1 | _T: ...
     def repeat(
         self, repeats: int | list[int], axis: AxisIndex | None = ...
     ) -> Series[S1]: ...

tests/test_series.py

@@ -16,6 +16,7 @@
     Any,
     Generic,
     TypeVar,
+    Union,
     cast,
 )
 
@@ -2875,6 +2876,24 @@ def test_round() -> None:
     check(assert_type(round(pd.Series([1], dtype=int)), "pd.Series[int]"), pd.Series)
 
 
+def test_get() -> None:
+    s_int = pd.Series([1, 2, 3], index=[1, 2, 3])
+
+    check(assert_type(s_int.get(1), Union[int, None]), np.int64)
+    check(assert_type(s_int.get(99), Union[int, None]), type(None))
+    check(assert_type(s_int.get(1, default=None), Union[int, None]), np.int64)
+    check(assert_type(s_int.get(1, default=2), int), np.int64)
+    check(assert_type(s_int.get(99, default="a"), Union[int, str]), str)
+
+    s_str = pd.Series(list("abc"), index=list("abc"))
+
+    check(assert_type(s_str.get("a"), Union[str, None]), str)
+    check(assert_type(s_str.get("z"), Union[str, None]), type(None))
+    check(assert_type(s_str.get("a", default=None), Union[str, None]), str)
+    check(assert_type(s_str.get("a", default="b"), str), str)
+    check(assert_type(s_str.get("z", default=True), Union[str, bool]), bool)
+
+
 def test_series_new_empty() -> None:
     # GH 826
     check(assert_type(pd.Series(), "pd.Series[Any]"), pd.Series)