如何切片 pandas DataFrame 以便 pylance 不会给出有关返回类型为 NDFrame 的错误？

Question

我有很大的 pandas DataFrame，我需要切片。

我注意到切片后 VSCode 假定 NDFrame 而不是 DataFrame。

我只使用二维数据帧，所以 NDFrame 让我感到困惑。我更喜欢只使用 DataFrame。 Python 代码按预期工作，这似乎是 pylance 扩展中的一个问题。

切片后如何接收DataFrame？

示例代码：

import pandas as pd
from pandas import DataFrame

df1: DataFrame = pd.DataFrame(
    {
        "A": ["A0", "A1", "A2", "A3"],
        "B": ["B0", "B1", "B2", "B3"],
        "C": ["C0", "C1", "C2", "C3"],
        "D": ["D0", "D1", "D2", "D3"],
    },
    index=[0, 1, 2, 3],
)

df2: DataFrame =df1[:2]

print(type(df1))
print(type(df2))

输出：

<class 'pandas.core.frame.DataFrame'>
<class 'pandas.core.frame.DataFrame'>

VSCode pylance类型检查错误：

Answer 1

要更正此问题，您可以使用 iloc:

    n = 2
    # iloc (note the difference for Series [:n] and DataFrame [:n,:]):
    df2: DataFrame = df.iloc[:n,:] # first n rows
    df3: DataFrame = df.iloc[-n:,:] # last n rows

或者简单地使用 head(n) 获取第一个 n 和 tail(n) 获取最后 n 行：

    n = 2
    df2: DataFrame = df.head(n)
    df3: DataFrame = df.tail(n)

以相同的方式 head(-n) 删除最后 n 行，tail(-n) 删除第 n 行：

    n = 2
    df2: DataFrame = df.head(-n)
    df3: DataFrame = df.tail(-n)

要了解更多信息，请阅读：https://pandas.pydata.org/pandas-docs/stable/user_guide/indexing.html

如何切片 pandas DataFrame 以便 pylance 不会给出有关返回类型为 NDFrame 的错误？

How to slice pandas DataFrame so pylance doesn't give error about returned type being NDFrame?

python

pandas

visual-studio-code

python-typing

pylance