Indexing Pandas MultiIndex and complying with mypy

Viewed 17

I am trying to get rid of some mypy errors when indexing into a pd.MultiIndex.
I installed pandas-stubs 1.4.4.220919 via pip install pandas-stubs.

Consider the following dataframe:

import numpy as np
import pandas as pd

col_level_1 = ["long_only", "short_only", "both"]
col_level_2 = ["min", "median", "mean", "max"]
idx = pd.bdate_range(start="2022-01-01", periods=5)
mi = pd.MultiIndex.from_product(
    (col_level_1, col_level_2), names=["direction", "statistic"]
)
df = pd.DataFrame(
    data=np.random.randint(low=0, high=10, size=(5, 12)), index=idx, columns=mi
)

print(df.loc[:, pd.IndexSlice[:, ["median", "mean"]]])

direction  long_only short_only   both long_only short_only both
statistic     median     median median      mean       mean mean
2022-01-03         6          0      4         6          2    2
2022-01-04         4          9      7         9          2    5
2022-01-05         8          3      0         8          6    4
2022-01-06         3          0      2         0          9    0
2022-01-07         3          8      4         0          5    3

The code is perfectly working fine. However, mypy throws an error:

Invalid index type "Tuple[slice, Tuple[slice, List[str]]]" for "_LocIndexerFrame"; expected type "Union[slice, ndarray[Any, dtype[signedinteger[_64Bit]]], Index, List[int], Series[int], Series[bool], ndarray[Any, dtype[bool_]], List[bool], List[<nothing>], Tuple[Union[slice, Index, List[int], Series[int], Series[bool], List[bool], List[<nothing>], Hashable], Union[List[<nothing>], slice, Series[bool], Callable[..., Any]]]]" [index]

0 Answers
Related