Pandas loc incorrect type with multi-index

[amgcc]

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

[Dr-Irv]

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:

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.

[amgcc]

@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]"

[amgcc]

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

[amgcc]

So the __getitem__ is part of _LocIndexer and the syntax is quite complicated but it does appear to be resolving incorrectly.

pandas-stubs/pandas-stubs/core/series.pyi

Line 169 in 8d5a827

class _LocIndexerSeries(_LocIndexer, Generic[S1]):

[Dr-Irv]

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

    @overload
    def __getitem__(
        self,
        idx: ScalarT
        | Callable[[DataFrame], ScalarT]

to

    @overload
    def __getitem__(self, idx: Scalar) -> Series | DataFrame: ...
    @overload
    def __getitem__(
        self,
        idx: Callable[[DataFrame], ScalarT]

and change the set of tests from

pandas-stubs/tests/test_frame.py

Line 1987 in 8d5a827

check(assert_type(df.loc[str_], pd.Series), pd.Series)

to

pandas-stubs/tests/test_frame.py

Line 1996 in 8d5a827

check(assert_type(df.loc[pd_timedelta], pd.Series), pd.Series)

to do assert_type() to check Series | DataFrame
and add tests from OP above and from #749 (comment)

PR welcome.