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

Issue:
Pandas loc incorrect type with multi-index
When using a multi-index dataframe, accessing values with loc returns another dataframe.  However, mypy seems to report this as an error expecting a series.  I extensively use multi-index dataframes and it makes usage of stubs very difficult.

**To Reproduce**
1. Provide a minimal runnable `pandas` example that is not properly checked by the stubs.
```
import pandas as pd

def main() -> None:
    idx = pd.MultiIndex.from_product(
        [["a", "b"], [1, 2]], names=["alpha", "num"]
    )

    df = pd.DataFrame({"a": range(4)}, index=idx)
    df2: pd.DataFrame = df.loc["a"]
    print(type(df))
    print(df)
    print(type(df2))
    print(df2)
    # reveal_locals()


if __name__ == "__main__":
    main()
```
3. Indicate which type checker you are using (`mypy` or  `pyright`).
mypy
4. Show the error message received from that type checker while checking your example.
```
demo.py:13: error: Incompatible types in assignment (expression has type "Series[Any]", variable has type "DataFrame")  [assignment]
```

**Please complete the following information:**
 - OS: [e.g. Windows, Linux, MacOS]
 Linux
 - OS Version [e.g. 22]
 CentOS Linux release 7.8.2003 (Core)
 - python version
 Python 3.8
 - version of type checker
 mypy 1.4.1 (compiled: yes)
 - version of installed `pandas-stubs`
pandas-stubs=2.0.1.230501=pyhd8ed1ab_0

**Additional context**
Add any other context about the problem here.

```
$ python demo.py 
<class 'pandas.core.frame.DataFrame'>
           a
alpha num   
a     1    0
      2    1
      3    2
b     1    3
      2    4
      3    5
c     1    6
      2    7
      3    8
<class 'pandas.core.frame.DataFrame'>
     a
num   
1    0
2    1
3    2

Process finished with exit code 0
```

Pandas loc incorrect type with multi-index
When using a multi-index dataframe, accessing values with loc returns another dataframe.  However, mypy seems to report this as an error expecting a series.  I extensively use multi-index dataframes and it makes usage of stubs very difficult.

**To Reproduce**
1. Provide a minimal runnable `pandas` example that is not properly checked by the stubs.
```
import pandas as pd

def main() -> None:
    idx = pd.MultiIndex.from_product(
        [["a", "b"], [1, 2]], names=["alpha", "num"]
    )

    df = pd.DataFrame({"a": range(4)}, index=idx)
    df2: pd.DataFrame = df.loc["a"]
    print(type(df))
    print(df)
    print(type(df2))
    print(df2)
    # reveal_locals()


if __name__ == "__main__":
    main()
```
3. Indicate which type checker you are using (`mypy` or  `pyright`).
mypy
4. Show the error message received from that type checker while checking your example.
```
demo.py:13: error: Incompatible types in assignment (expression has type "Series[Any]", variable has type "DataFrame")  [assignment]
```

**Please complete the following information:**
 - OS: [e.g. Windows, Linux, MacOS]
 Linux
 - OS Version [e.g. 22]
 CentOS Linux release 7.8.2003 (Core)
 - python version
 Python 3.8
 - version of type checker
 mypy 1.4.1 (compiled: yes)
 - version of installed `pandas-stubs`
pandas-stubs=2.0.1.230501=pyhd8ed1ab_0

**Additional context**
Add any other context about the problem here.

```
$ python demo.py 
<class 'pandas.core.frame.DataFrame'>
           a
alpha num   
a     1    0
      2    1
      3    2
b     1    3
      2    4
      3    5
c     1    6
      2    7
      3    8
<class 'pandas.core.frame.DataFrame'>
     a
num   
1    0
2    1
3    2

Process finished with exit code 0
```

Pandas loc incorrect type with multi-index
When using a multi-index dataframe, accessing values with loc returns another dataframe.  However, mypy seems to report this as an error expecting a series.  I extensively use multi-index dataframes and it makes usage of stubs very difficult.

**To Reproduce**
1. Provide a minimal runnable `pandas` example that is not properly checked by the stubs.
```
import pandas as pd

def main() -> None:
    idx = pd.MultiIndex.from_product(
        [["a", "b"], [1, 2]], names=["alpha", "num"]
    )

    df = pd.DataFrame({"a": range(4)}, index=idx)
    df2: pd.DataFrame = df.loc["a"]
    print(type(df))
    print(df)
    print(type(df2))
    print(df2)
    # reveal_locals()


if __name__ == "__main__":
    main()
```
3. Indicate which type checker you are using (`mypy` or  `pyright`).
mypy
4. Show the error message received from that type checker while checking your example.
```
demo.py:13: error: Incompatible types in assignment (expression has type "Series[Any]", variable has type "DataFrame")  [assignment]
```

**Please complete the following information:**
 - OS: [e.g. Windows, Linux, MacOS]
 Linux
 - OS Version [e.g. 22]
 CentOS Linux release 7.8.2003 (Core)
 - python version
 Python 3.8
 - version of type checker
 mypy 1.4.1 (compiled: yes)
 - version of installed `pandas-stubs`
pandas-stubs=2.0.1.230501=pyhd8ed1ab_0

**Additional context**
Add any other context about the problem here.

```
$ python demo.py 
<class 'pandas.core.frame.DataFrame'>
           a
alpha num   
a     1    0
      2    1
      3    2
b     1    3
      2    4
      3    5
c     1    6
      2    7
      3    8
<class 'pandas.core.frame.DataFrame'>
     a
num   
1    0
2    1
3    2

Process finished with exit code 0
```



Discussion / hints:
The problem with your example is that the expression `df.loc["a"]` is ambiguous from a static typing perspective.  It's really not possible in a static typing context to know what `.loc` will do, since the type of the underlying index backing the `DataFrame` is not known via static analysis. 

For example:
```python
df3 = pd.DataFrame({"x": range(2)}, index=pd.Index(["a", "b"]))
s4 = df3.loc["a"]
```
Here, the runtime type of `s4` is `Series`.

So you would have to do a `cast` if you want to use `.loc` in this way.  Right now, we return `Series | DataFrame` for this operation, as the 2 examples illustrate that you can get either type.




@Dr-Irv Thanks for the quick reply.  It makes perfect sense for pandas-stubs to return Dataframe or Series from the `loc` operation.  However that isn't what is happening in my case.  If you look closely at my submission, you will see that mypy is resolving this instead as a Series causing my issue.  

Using `reveal_type`:

```
reveal_type(df2)
demo.py:18: note: Revealed type is "pandas.core.series.Series[Any]"
```
What you described as the return type matches what I find in `pytables.pyi` but mypy resolves it as a Series instead.  My understanding of pyi syntax is very limited so maybe there is some explanation of this behavior.

```
def __getitem__(self, key: str) -> DataFrame | Series: ...
```

Note this is the `__getitem__` in `pd.DataFrame` and not on `df.loc`
So the `__getitem__` is part of `_LocIndexer` and the syntax is quite complicated but it does appear to be resolving incorrectly. 

https://github.com/pandas-dev/pandas-stubs/blob/8d5a8279d448afa188f5160e7fe787c1fcb6015e/pandas-stubs/core/series.pyi#L169
> It makes perfect sense for pandas-stubs to return Dataframe or Series from the `loc` operation. However that isn't what is happening in my case. If you look closely at my submission, you will see that mypy is resolving this instead as a Series causing my issue.

Yes, I see that now.  Reopening.

In `core/frame.pyi`, should change at
https://github.com/pandas-dev/pandas-stubs/blob/8d5a8279d448afa188f5160e7fe787c1fcb6015e/pandas-stubs/core/frame.pyi#L188C1-L191C21

```python
    @overload
    def __getitem__(
        self,
        idx: ScalarT
        | Callable[[DataFrame], ScalarT]
```
to
```python
    @overload
    def __getitem__(self, idx: Scalar) -> Series | DataFrame: ...
    @overload
    def __getitem__(
        self,
        idx: Callable[[DataFrame], ScalarT]
```
and change the set of tests from
https://github.com/pandas-dev/pandas-stubs/blob/8d5a8279d448afa188f5160e7fe787c1fcb6015e/tests/test_frame.py#L1987
to
https://github.com/pandas-dev/pandas-stubs/blob/8d5a8279d448afa188f5160e7fe787c1fcb6015e/tests/test_frame.py#L1996
to do `assert_type()` to check `Series | DataFrame` 
and add tests from OP above and from https://github.com/pandas-dev/pandas-stubs/issues/749#issuecomment-1640816583

PR welcome.

The problem with your example is that the expression `df.loc["a"]` is ambiguous from a static typing perspective.  It's really not possible in a static typing context to know what `.loc` will do, since the type of the underlying index backing the `DataFrame` is not known via static analysis. 

For example:
```python
df3 = pd.DataFrame({"x": range(2)}, index=pd.Index(["a", "b"]))
s4 = df3.loc["a"]
```
Here, the runtime type of `s4` is `Series`.

So you would have to do a `cast` if you want to use `.loc` in this way.  Right now, we return `Series | DataFrame` for this operation, as the 2 examples illustrate that you can get either type.




@Dr-Irv Thanks for the quick reply.  It makes perfect sense for pandas-stubs to return Dataframe or Series from the `loc` operation.  However that isn't what is happening in my case.  If you look closely at my submission, you will see that mypy is resolving this instead as a Series causing my issue.  

Using `reveal_type`:

```
reveal_type(df2)
demo.py:18: note: Revealed type is "pandas.core.series.Series[Any]"
```
What you described as the return type matches what I find in `pytables.pyi` but mypy resolves it as a Series instead.  My understanding of pyi syntax is very limited so maybe there is some explanation of this behavior.

```
def __getitem__(self, key: str) -> DataFrame | Series: ...
```

Note this is the `__getitem__` in `pd.DataFrame` and not on `df.loc`
So the `__getitem__` is part of `_LocIndexer` and the syntax is quite complicated but it does appear to be resolving incorrectly. 

https://github.com/pandas-dev/pandas-stubs/blob/8d5a8279d448afa188f5160e7fe787c1fcb6015e/pandas-stubs/core/series.pyi#L169
> It makes perfect sense for pandas-stubs to return Dataframe or Series from the `loc` operation. However that isn't what is happening in my case. If you look closely at my submission, you will see that mypy is resolving this instead as a Series causing my issue.

Yes, I see that now.  Reopening.

In `core/frame.pyi`, should change at
https://github.com/pandas-dev/pandas-stubs/blob/8d5a8279d448afa188f5160e7fe787c1fcb6015e/pandas-stubs/core/frame.pyi#L188C1-L191C21

```python
    @overload
    def __getitem__(
        self,
        idx: ScalarT
        | Callable[[DataFrame], ScalarT]
```
to
```python
    @overload
    def __getitem__(self, idx: Scalar) -> Series | DataFrame: ...
    @overload
    def __getitem__(
        self,
        idx: Callable[[DataFrame], ScalarT]
```
and change the set of tests from
https://github.com/pandas-dev/pandas-stubs/blob/8d5a8279d448afa188f5160e7fe787c1fcb6015e/tests/test_frame.py#L1987
to
https://github.com/pandas-dev/pandas-stubs/blob/8d5a8279d448afa188f5160e7fe787c1fcb6015e/tests/test_frame.py#L1996
to do `assert_type()` to check `Series | DataFrame` 
and add tests from OP above and from https://github.com/pandas-dev/pandas-stubs/issues/749#issuecomment-1640816583

PR welcome.

The problem with your example is that the expression `df.loc["a"]` is ambiguous from a static typing perspective.  It's really not possible in a static typing context to know what `.loc` will do, since the type of the underlying index backing the `DataFrame` is not known via static analysis. 

For example:
```python
df3 = pd.DataFrame({"x": range(2)}, index=pd.Index(["a", "b"]))
s4 = df3.loc["a"]
```
Here, the runtime type of `s4` is `Series`.

So you would have to do a `cast` if you want to use `.loc` in this way.  Right now, we return `Series | DataFrame` for this operation, as the 2 examples illustrate that you can get either type.




@Dr-Irv Thanks for the quick reply.  It makes perfect sense for pandas-stubs to return Dataframe or Series from the `loc` operation.  However that isn't what is happening in my case.  If you look closely at my submission, you will see that mypy is resolving this instead as a Series causing my issue.  

Using `reveal_type`:

```
reveal_type(df2)
demo.py:18: note: Revealed type is "pandas.core.series.Series[Any]"
```
What you described as the return type matches what I find in `pytables.pyi` but mypy resolves it as a Series instead.  My understanding of pyi syntax is very limited so maybe there is some explanation of this behavior.

```
def __getitem__(self, key: str) -> DataFrame | Series: ...
```

Note this is the `__getitem__` in `pd.DataFrame` and not on `df.loc`
So the `__getitem__` is part of `_LocIndexer` and the syntax is quite complicated but it does appear to be resolving incorrectly. 

https://github.com/pandas-dev/pandas-stubs/blob/8d5a8279d448afa188f5160e7fe787c1fcb6015e/pandas-stubs/core/series.pyi#L169
> It makes perfect sense for pandas-stubs to return Dataframe or Series from the `loc` operation. However that isn't what is happening in my case. If you look closely at my submission, you will see that mypy is resolving this instead as a Series causing my issue.

Yes, I see that now.  Reopening.

In `core/frame.pyi`, should change at
https://github.com/pandas-dev/pandas-stubs/blob/8d5a8279d448afa188f5160e7fe787c1fcb6015e/pandas-stubs/core/frame.pyi#L188C1-L191C21

```python
    @overload
    def __getitem__(
        self,
        idx: ScalarT
        | Callable[[DataFrame], ScalarT]
```
to
```python
    @overload
    def __getitem__(self, idx: Scalar) -> Series | DataFrame: ...
    @overload
    def __getitem__(
        self,
        idx: Callable[[DataFrame], ScalarT]
```
and change the set of tests from
https://github.com/pandas-dev/pandas-stubs/blob/8d5a8279d448afa188f5160e7fe787c1fcb6015e/tests/test_frame.py#L1987
to
https://github.com/pandas-dev/pandas-stubs/blob/8d5a8279d448afa188f5160e7fe787c1fcb6015e/tests/test_frame.py#L1996
to do `assert_type()` to check `Series | DataFrame` 
and add tests from OP above and from https://github.com/pandas-dev/pandas-stubs/issues/749#issuecomment-1640816583

PR welcome.


Fix patch:
diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi
index 45f02ca2..116ba52e 100644
--- a/pandas-stubs/core/frame.pyi
+++ b/pandas-stubs/core/frame.pyi
@@ -156,6 +156,8 @@ class _iLocIndexerFrame(_iLocIndexer):
     ) -> None: ...
 
 class _LocIndexerFrame(_LocIndexer):
+    @overload
+    def __getitem__(self, idx: Scalar) -> Series | DataFrame: ...
     @overload
     def __getitem__(
         self,
@@ -184,8 +186,7 @@ class _LocIndexerFrame(_LocIndexer):
     @overload
     def __getitem__(
         self,
-        idx: ScalarT
-        | Callable[[DataFrame], ScalarT]
+        idx: Callable[[DataFrame], ScalarT]
         | tuple[
             IndexType
             | MaskType


Test patch:
diff --git a/tests/test_frame.py b/tests/test_frame.py
index cd4ace1f..c193f2b4 100644
--- a/tests/test_frame.py
+++ b/tests/test_frame.py
@@ -2199,16 +2199,16 @@ def test_frame_scalars_slice() -> None:
 
     # Note: bool_ cannot be tested since the index is object and pandas does not
     # support boolean access using loc except when the index is boolean
-    check(assert_type(df.loc[str_], pd.Series), pd.Series)
-    check(assert_type(df.loc[bytes_], pd.Series), pd.Series)
-    check(assert_type(df.loc[date], pd.Series), pd.Series)
-    check(assert_type(df.loc[datetime_], pd.Series), pd.Series)
-    check(assert_type(df.loc[timedelta], pd.Series), pd.Series)
-    check(assert_type(df.loc[int_], pd.Series), pd.Series)
-    check(assert_type(df.loc[float_], pd.Series), pd.Series)
-    check(assert_type(df.loc[complex_], pd.Series), pd.Series)
-    check(assert_type(df.loc[timestamp], pd.Series), pd.Series)
-    check(assert_type(df.loc[pd_timedelta], pd.Series), pd.Series)
+    check(assert_type(df.loc[str_], Union[pd.Series, pd.DataFrame]), pd.Series)
+    check(assert_type(df.loc[bytes_], Union[pd.Series, pd.DataFrame]), pd.Series)
+    check(assert_type(df.loc[date], Union[pd.Series, pd.DataFrame]), pd.Series)
+    check(assert_type(df.loc[datetime_], Union[pd.Series, pd.DataFrame]), pd.Series)
+    check(assert_type(df.loc[timedelta], Union[pd.Series, pd.DataFrame]), pd.Series)
+    check(assert_type(df.loc[int_], Union[pd.Series, pd.DataFrame]), pd.Series)
+    check(assert_type(df.loc[float_], Union[pd.Series, pd.DataFrame]), pd.Series)
+    check(assert_type(df.loc[complex_], Union[pd.Series, pd.DataFrame]), pd.Series)
+    check(assert_type(df.loc[timestamp], Union[pd.Series, pd.DataFrame]), pd.Series)
+    check(assert_type(df.loc[pd_timedelta], Union[pd.Series, pd.DataFrame]), pd.Series)
     check(assert_type(df.loc[none], pd.Series), pd.Series)
 
     check(assert_type(df.loc[:, str_], pd.Series), pd.Series)
@@ -2223,11 +2223,20 @@ def test_frame_scalars_slice() -> None:
     check(assert_type(df.loc[:, pd_timedelta], pd.Series), pd.Series)
     check(assert_type(df.loc[:, none], pd.Series), pd.Series)
 
+    # GH749
+
+    multi_idx = pd.MultiIndex.from_product([["a", "b"], [1, 2]], names=["alpha", "num"])
+    df2 = pd.DataFrame({"col1": range(4)}, index=multi_idx)
+    check(assert_type(df2.loc[str_], Union[pd.Series, pd.DataFrame]), pd.DataFrame)
+
+    df3 = pd.DataFrame({"x": range(2)}, index=pd.Index(["a", "b"]))
+    check(assert_type(df3.loc[str_], Union[pd.Series, pd.DataFrame]), pd.Series)
+
 
 def test_boolean_loc() -> None:
     # Booleans can only be used in loc when the index is boolean
     df = pd.DataFrame([[0, 1], [1, 0]], columns=[True, False], index=[True, False])
-    check(assert_type(df.loc[True], pd.Series), pd.Series)
+    check(assert_type(df.loc[True], Union[pd.Series, pd.DataFrame]), pd.Series)
     check(assert_type(df.loc[:, False], pd.Series), pd.Series)