python - 键入具有特定列的提示 pandas 数据框

Question

假设我有以下功能:

def foo(df: pd.DataFrame) -> pd.DataFrame:
    x = df["x"]
    y = df["y"]
    df["xy"] = x * y

    return df

有没有一种方法可以暗示我的函数正在接受一个必须有“x”和“y”列的数据框，并且它将返回一个包含“x”、“y”和“xy”的数据框"列，而不仅仅是一般的数据框？

Answer 1

好的，所以，我不确定这是否是正确的实现方式，但似乎对我有用。如果您发现任何错误或替代方案，请告诉我，我可以编辑回复，但我的解决方案基本上是创建一个新类并实现 __class_getitem__ 方法，如 Pep 560 中所示，这是我的最终代码:

from typing import List
import pandas as pd


GenericAlias = type(List[str])

class MyDataFrame(pd.DataFrame):
    __class_getitem__ = classmethod(GenericAlias)


def foo(df: MyDataFrame[["x", "y"]]) -> MyDataFrame[["x", "y", "xy"]]:
    df["xy"] = df["x"] * df["y"]
    return df