如何为含指定列的Pandas DataFrame添加类型提示

可行实现方案有以下三种，可根据你的项目依赖版本、校验需求选择： #

方案1：Pandas 原生类型标注（Pandas 2.0+ 适用） #

Pandas 2.0 及以上版本原生支持泛型 DataFrame 类型，配合 pandas-stubs 可以让 mypy、pyright 等静态类型检查器识别列结构约束：

import pandas as pd
from pandas.api.typing import DataFrame

# 定义入参 DataFrame 类型：包含 x、y 两列，类型为64位浮点数
InputDF = DataFrame[
    [("x", pd.Float64Dtype), ("y", pd.Float64Dtype)]
]
# 定义返回值 DataFrame 类型：额外新增 xy 列
OutputDF = DataFrame[
    [("x", pd.Float64Dtype), ("y", pd.Float64Dtype), ("xy", pd.Float64Dtype)]
]

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

特点：无额外第三方依赖（仅需安装开发依赖 pandas-stubs），静态检查覆盖完整，运行时无额外性能损耗。

方案2：Pandera 运行时+静态双校验 #

如果需要运行时也自动校验列结构和字段类型，可以用专门面向 Pandas 的校验库 pandera：

import pandas as pd
import pandera as pa
from pandera.typing import DataFrame, Series

# 定义入参结构约束
class InputSchema(pa.DataFrameModel):
    x: Series[float]
    y: Series[float]

# 继承入参结构，新增返回值要求的 xy 列
class OutputSchema(InputSchema):
    xy: Series[float]

# 装饰器开启运行时自动校验
@pa.check_types
def foo(df: DataFrame[InputSchema]) -> DataFrame[OutputSchema]:
    x = df["x"]
    y = df["y"]
    df["xy"] = x * y
    return df

特点：除了静态类型提示，传入数据不符合结构时会直接抛出清晰的校验错误，避免后续取列、计算时才触发难以定位的报错，适合对数据正确性要求高的场景。

方案3：Annotated 标记（低版本 Pandas 兼容方案） #

如果使用的 Pandas 版本低于 2.0，无法使用原生泛型标注，可以用 typing.Annotated 给通用 pd.DataFrame 增加显式结构标记：

from typing import Annotated
import pandas as pd

# 用字符串标注结构约束，供开发者和工具识别
InputDF = Annotated[pd.DataFrame, "必须包含x、y两列，均为数值类型"]
OutputDF = Annotated[pd.DataFrame, "必须包含x、y、xy三列，均为数值类型"]

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

特点：兼容性最强，所有 Python 3.9+ 版本都支持，仅作为显式的结构说明，无强制静态/运行时校验，适合临时项目或者依赖版本受限的场景。

内容的提问来源于stack exchange，提问作者João Areias