如何输入多索引熊猫数据框 Python3.8+ 的提示？

Posted 2023-03-12

【中文标题】如何输入多索引熊猫数据框 Python3.8+ 的提示？

【英文标题】：How to type hinting of multi-indexed pandas dataframe, Python3.8+?

【发布时间】：2021-12-01 08:44:18

【问题描述】：

数据框df的列是多索引的。

>>>df
    c1      c2    
    d1  d2  d1  d2
r1  11  12  21  22

因此，使用代码list(df.columns.levels[0] 创建list 的['c1', 'c2']。以下代码可以正常运行。但是 PyCharm 在print(list(df.columns.levels[0])) 行中将levels 标记为警告Unresolved attribute reference 'levels' for class 'Index'。

我认为这可以通过print_columns 的函数签名中更具体的类型提示来解决。在下面的代码中，我只使用了df: pd.DataFrame 类型。 PyCharm 认为它是一个没有多索引的简单数据框。

你能告诉我如何改进类型提示吗？或者您有其他解决方案可以从 PyCharm 中删除警告吗？

谢谢。

import pandas as pd


def get_df() -> pd.DataFrame:
    data = 
        ('c1', 'd1'): [11],
        ('c1', 'd2'): [12],
        ('c2', 'd1'): [21],
        ('c2', 'd2'): [22],
    

    df = pd.DataFrame.from_dict(data)
    df.index = ['r1']
    return df


def print_columns(df: pd.DataFrame):
    print(list(df.columns.levels[0]))  # prints ['c1', 'c2']


if __name__ == '__main__':
    df_ = get_df()
    print_columns(df=df_)

【参考方案1】：

根据Pandas documentation，可以像这样对多级对象进行类型提示：

import pandas as pd
...
def print_columns(df: pd.MultiIndex) -> None:
    print(list(df.columns.levels[0]))

【讨论】：

感谢您的回答。但是对于df: pd.MultiIndex，PyCharm 将columns 标记为警告而不是levels。

我担心这可能是 PyCharm 的限制，根据这个类似的帖子 ***.com/a/42440466/11246056。但这篇其他帖子可能会有所帮助：***.com/a/67006481/11246056