如何为返回Pandas DataFrame的函数添加含列信息的类型提示？

下面是几种能在VS Code、PyCharm和Jupyter Notebook中良好集成的方案，满足你对DataFrame列类型提示的需求：

1. 基于pandas-stubs的TypedDict方案（推荐轻量静态提示） #

先安装类型注解依赖：

pip install pandas-stubs

然后通过TypedDict定义列的结构，再结合pd.DataFrame的泛型注解：

from typing import TypedDict
import pandas as pd

# 定义DataFrame的列与对应类型
class GeneratedData(TypedDict):
    a: int
    b: float
    c: str
    d: float  # 美元金额类型

def generate_data(bunch, of, inputs) -> pd.DataFrame[GeneratedData]:
    """Massages the input to a nice and easy DataFrame.
    
    :return: DataFrame with specified columns and types
    """
    # 函数实现逻辑
    ...

主流编辑器会识别TypedDict的定义，提供列名补全和类型提示，Jupyter Notebook中也能通过代码提示看到列信息。

2. pydantic-dataframe方案（带运行时校验） #

如果需要同时做静态提示和运行时数据校验，可使用这个库：

pip install pydantic-dataframe

用Pydantic模型定义列规则，再包装成DataFrame类型：

from pydantic import BaseModel
from pydantic_dataframe import DataFrameModel

# 定义列的类型规范
class DataSchema(BaseModel):
    a: int
    b: float
    c: str
    d: float

# 生成对应的DataFrame类型
class GeneratedDataFrame(DataFrameModel[DataSchema]):
    pass

def generate_data(bunch, of, inputs) -> GeneratedDataFrame:
    """Massages the input to a nice and easy DataFrame.
    
    :return: DataFrame matching the specified column types
    """
    df = ...  # 生成DataFrame的逻辑
    return GeneratedDataFrame(df)

这个方案会在运行时自动校验返回的DataFrame是否符合列规则，同时编辑器能识别类型并提供提示。

3. 文档字符串增强方案（零依赖） #

如果不想引入额外包，可优化文档字符串的格式，让编辑器自动解析列信息：

import pandas as pd

def generate_data(bunch, of, inputs) -> pd.DataFrame:
    """Massages the input to a nice and easy DataFrame.
    
    :return: pd.DataFrame with columns:
        - a: int
        - b: float
        - c: str
        - d: float (US dollars)
    """
    ...

VS Code、PyCharm会解析文档字符串中的列说明，用户悬停函数时就能看到详细的列类型信息，Jupyter Notebook中也能通过help(generate_data)查看完整说明。

内容的提问来源于stack exchange，提问作者Mikko Ohtamaa