使用特定列键入提示 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

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

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 

编辑

Pandera库是一个更好的选择，当我发布问题时我并不知道。

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