Add get method to DataFrame

Author: skatsuta

pandas-stubs/core/frame.pyi

@@ -13,6 +13,7 @@ from typing import (
     Any,
     ClassVar,
     Literal,
+    TypeVar,
     overload,
 )
 
@@ -127,6 +128,8 @@ from pandas.plotting import PlotAccessor
 _str = str
 _bool = bool
 
+_T = TypeVar("_T")
+
 class _iLocIndexerFrame(_iLocIndexer):
     @overload
     def __getitem__(self, idx: tuple[int, int]) -> Scalar: ...
@@ -1683,7 +1686,14 @@ class DataFrame(NDFrame, OpsMixin):
     # def from_dict
     # def from_records
     def ge(self, other, axis: Axis = ..., level: Level | None = ...) -> DataFrame: ...
-    # def get
+    @overload
+    def get(self, key: Hashable, default: None = ...) -> Series | None: ...
+    @overload
+    def get(self, key: Hashable, default: _T = ...) -> Series | _T: ...
+    @overload
+    def get(self, key: list[Hashable], default: None = ...) -> DataFrame | None: ...
+    @overload
+    def get(self, key: list[Hashable], default: _T = ...) -> DataFrame | _T: ...
     def gt(self, other, axis: Axis = ..., level: Level | None = ...) -> DataFrame: ...
     def head(self, n: int = ...) -> DataFrame: ...
     def infer_objects(self) -> DataFrame: ...

tests/test_frame.py

@@ -3107,3 +3107,33 @@ def test_itertuples() -> None:
     for item in df.itertuples():
         check(assert_type(item, _PandasNamedTuple), tuple)
         assert_type(item.a, Scalar)
+
+
+def test_get() -> None:
+    df = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6], "c": [7, 8, 9]})
+
+    check(assert_type(df.get("a"), Union[pd.Series, None]), pd.Series, np.int64)
+    check(assert_type(df.get("z"), Union[pd.Series, None]), type(None))
+    check(
+        assert_type(df.get("a", default=None), Union[pd.Series, None]),
+        pd.Series,
+        np.int64,
+    )
+    check(
+        assert_type(df.get("a", default=1), Union[pd.Series, int]), pd.Series, np.int64
+    )
+    check(assert_type(df.get("z", default=1), Union[pd.Series, int]), int)
+
+    check(assert_type(df.get([]), Union[pd.DataFrame, None]), pd.DataFrame)
+    check(assert_type(df.get(["a"]), Union[pd.DataFrame, None]), pd.DataFrame)
+    check(assert_type(df.get(["a", "b"]), Union[pd.DataFrame, None]), pd.DataFrame)
+    check(assert_type(df.get(["z"]), Union[pd.DataFrame, None]), type(None))
+    check(
+        assert_type(df.get(["a", "b"], default=None), Union[pd.DataFrame, None]),
+        pd.DataFrame,
+    )
+    check(
+        assert_type(df.get(["a", "b"], default=1), Union[pd.DataFrame, int]),
+        pd.DataFrame,
+    )
+    check(assert_type(df.get(["z"], default=1), Union[pd.DataFrame, int]), int)