Repository: pandas-dev/pandas-stubs (language: python)

Issue:
pd.Series.get causes typing error in second argument (default value)
**Describe the bug**
`pd.Series.get()` seems to ignore the Series' type and causes an unexpected mypy error in the default argument.

**To Reproduce**
Minimal runnable `pandas` example:
```python
from typing import cast
import pandas as pd

s = pd.Series([1, 3, 3, 7, "l33t"], index=["a", "b", "c", "d"])

i = s.get("b", -1)

j = cast(pd.Series[int], s).get("b", -1)
```

I am using `mypy`, and the error message I receive is
`Argument 2 to "get" of "NDFrame" has incompatible type "int"; expected "ExtensionDtype | str | dtype[generic] | type[str] | type[complex] | type[bool] | type[object] | None"  [arg-type]mypy(error)`

![image](https://github.com/pandas-dev/pandas-stubs/assets/88887073/1ccc6a48-72d4-4053-8985-daf116450a3b)

I can understand why the `i` line can cause problems (as the Series' dtype is not uniform), but at least the `j` line should work fine in my opinion, and I have no idea how to fix this other than via `# type: ignore`.

**Please complete the following information:**
 - OS: Windows 10
 - Python 3.10.13
 - `mypy=1.6.1` (shipped with VS Code extension; env contains `1.5.1`)
 - `pandas-stubs==2.1.1.230928` (via `conda-forge`)

**Additional info:**
`get`'s signature:
```
(method) def get(
    key: object,
    default: Dtype | None = ...
) -> Dtype
```
pd.Series.get causes typing error in second argument (default value)
**Describe the bug**
`pd.Series.get()` seems to ignore the Series' type and causes an unexpected mypy error in the default argument.

**To Reproduce**
Minimal runnable `pandas` example:
```python
from typing import cast
import pandas as pd

s = pd.Series([1, 3, 3, 7, "l33t"], index=["a", "b", "c", "d"])

i = s.get("b", -1)

j = cast(pd.Series[int], s).get("b", -1)
```

I am using `mypy`, and the error message I receive is
`Argument 2 to "get" of "NDFrame" has incompatible type "int"; expected "ExtensionDtype | str | dtype[generic] | type[str] | type[complex] | type[bool] | type[object] | None"  [arg-type]mypy(error)`

![image](https://github.com/pandas-dev/pandas-stubs/assets/88887073/1ccc6a48-72d4-4053-8985-daf116450a3b)

I can understand why the `i` line can cause problems (as the Series' dtype is not uniform), but at least the `j` line should work fine in my opinion, and I have no idea how to fix this other than via `# type: ignore`.

**Please complete the following information:**
 - OS: Windows 10
 - Python 3.10.13
 - `mypy=1.6.1` (shipped with VS Code extension; env contains `1.5.1`)
 - `pandas-stubs==2.1.1.230928` (via `conda-forge`)

**Additional info:**
`get`'s signature:
```
(method) def get(
    key: object,
    default: Dtype | None = ...
) -> Dtype
```


Discussion / hints:
Thanks for the report.  The definition of `get()` in `core/generic.pyi` is totally wrong!  It currently reads:
```python
    def get(self, key: object, default: Dtype | None = ...) -> Dtype: ...
```
The fixes that should improve it include:

1. Have different versions of `get()` for `DataFrame` and `Series` in `core/frame.pyi` and `core/series.pyi` respectively, removing from `core/generic.pyi`
2. For the `DataFrame` version, need to distinguish between single keys that return a series and a sequence that returns a `DataFrame` via overloads
3. For the `Series` case, the `key` should be a `Hashable`, `default` should have type `S1`, and return value should be `S1`
4. Appropriate tests for both cases.


Thanks for the report.  The definition of `get()` in `core/generic.pyi` is totally wrong!  It currently reads:
```python
    def get(self, key: object, default: Dtype | None = ...) -> Dtype: ...
```
The fixes that should improve it include:

1. Have different versions of `get()` for `DataFrame` and `Series` in `core/frame.pyi` and `core/series.pyi` respectively, removing from `core/generic.pyi`
2. For the `DataFrame` version, need to distinguish between single keys that return a series and a sequence that returns a `DataFrame` via overloads
3. For the `Series` case, the `key` should be a `Hashable`, `default` should have type `S1`, and return value should be `S1`
4. Appropriate tests for both cases.



Fix patch:
diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi
index faae3599..8656099c 100644
--- a/pandas-stubs/core/frame.pyi
+++ b/pandas-stubs/core/frame.pyi
@@ -114,6 +114,7 @@ from pandas._typing import (
     StorageOptions,
     StrLike,
     Suffixes,
+    T as _T,
     TimestampConvention,
     ValidationOptions,
     WriteBuffer,
@@ -1696,7 +1697,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: ...
diff --git a/pandas-stubs/core/generic.pyi b/pandas-stubs/core/generic.pyi
index 431b9dfe..e1032a1d 100644
--- a/pandas-stubs/core/generic.pyi
+++ b/pandas-stubs/core/generic.pyi
@@ -34,7 +34,6 @@ from pandas._typing import (
     AxisIndex,
     CompressionOptions,
     CSVQuoting,
-    Dtype,
     DtypeArg,
     DtypeBackend,
     FilePath,
@@ -299,7 +298,6 @@ class NDFrame(indexing.IndexingMixin):
         self, indices, axis=..., is_copy: _bool | None = ..., **kwargs
     ) -> Self: ...
     def __delitem__(self, idx: Hashable) -> None: ...
-    def get(self, key: object, default: Dtype | None = ...) -> Dtype: ...
     def reindex_like(
         self,
         other,
diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi
index 383c8df9..d8c691a3 100644
--- a/pandas-stubs/core/series.pyi
+++ b/pandas-stubs/core/series.pyi
@@ -141,6 +141,7 @@ from pandas._typing import (
     SortKind,
     StrDtypeArg,
     StrLike,
+    T,
     TimedeltaDtypeArg,
     TimestampConvention,
     TimestampDtypeArg,
@@ -381,6 +382,12 @@ 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) -> S1: ...
+    @overload
+    def get(self, key: Hashable, default: T) -> S1 | T: ...
     def repeat(
         self, repeats: int | list[int], axis: AxisIndex | None = ...
     ) -> Series[S1]: ...


Test patch:
diff --git a/tests/test_frame.py b/tests/test_frame.py
index ba4ded9a..96ee6d52 100644
--- a/tests/test_frame.py
+++ b/tests/test_frame.py
@@ -3107,3 +3107,38 @@ 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]})
+
+    # Get single column
+    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("z", default=None), Union[pd.Series, None]), type(None))
+    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)
+
+    # Get multiple columns
+    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(["z"], default=None), Union[pd.DataFrame, None]), type(None)
+    )
+    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)
diff --git a/tests/test_series.py b/tests/test_series.py
index 96e644a2..e12dba47 100644
--- a/tests/test_series.py
+++ b/tests/test_series.py
@@ -16,6 +16,7 @@
     Any,
     Generic,
     TypeVar,
+    Union,
     cast,
 )
 
@@ -2875,6 +2876,26 @@ 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(99, default=None), Union[int, None]), type(None))
+    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("z", default=None), Union[str, None]), type(None))
+    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)