如何为Polars DataFrame实现类似泛型的Python类型提示？

由于Polars的pl.DataFrame并非泛型类型，无法直接用DataFrame[MySchema]的写法做类型提示，以下是几种实用的实现方式：

1. 泛型子类+类型别名（推荐） #

通过Python的Generic和TypeVar封装泛型DataFrame，结合TypedDict定义schema，实现类似泛型的类型提示效果：

from typing import TypeAlias, TypeVar, Generic, TypedDict
import polars as pl

# 定义Schema的TypeVar，绑定到TypedDict
SchemaT = TypeVar("SchemaT", bound=TypedDict)

# 创建泛型DataFrame子类
class TypedDataFrame(pl.DataFrame, Generic[SchemaT]):
    pass

# 用TypedDict定义具体的schema
class MyFrameSchema(TypedDict):
    price: float
    quantity: int
    product_name: str

# 定义自定义类型别名
TMyDataFrame: TypeAlias = TypedDataFrame[MyFrameSchema]

# 使用示例
def calculate_total(df: TMyDataFrame) -> TMyDataFrame:
    return df.with_columns(total=pl.col("price") * pl.col("quantity"))

这种方式能让IDE准确识别列的类型，类型检查器（如mypy）也能正常工作。

2. 结合Annotated附加schema信息 #

利用Annotated给pl.DataFrame附加schema元数据，适合不需要修改DataFrame结构的场景：

from typing import TypeAlias, Annotated, TypedDict
import polars as pl

class MyFrameSchema(TypedDict):
    price: float
    quantity: int
    product_name: str

# 定义带schema注解的类型别名
TMyDataFrame: TypeAlias = Annotated[pl.DataFrame, MyFrameSchema]

# 使用示例
def filter_expensive(df: TMyDataFrame) -> TMyDataFrame:
    return df.filter(pl.col("price") > 100)

注意：部分类型检查器（如mypy）需要额外插件支持解析Annotated中的schema信息，但主流IDE（PyCharm、VS Code）能提供基础的列名提示。

3. 直接子类化pl.DataFrame #

如果需要给DataFrame添加自定义方法，同时做类型提示，可以直接子类化：

import polars as pl
from typing import TypeAlias

class MyDataFrame(pl.DataFrame):
    def get_total(self) -> pl.Series:
        # 直接在方法里做类型约束
        return self["price"] * self["quantity"]

# 定义类型别名
TMyDataFrame: TypeAlias = MyDataFrame

# 使用示例
def process_data(df: TMyDataFrame) -> TMyDataFrame:
    return df.with_columns(total=df.get_total())

这种方式需要手动维护列的类型逻辑，适合schema固定且需要扩展方法的场景。

关于SchemaDict的用法 #

Polars的SchemaDict是dict[str, PolarsDataType | type]的类型别名，主要用于标注构造DataFrame时的schema参数，而非DataFrame本身的类型：

from polars.type_aliases import SchemaDict
import polars as pl

def create_order_df(schema: SchemaDict) -> pl.DataFrame:
    return pl.DataFrame(schema=schema)

# 定义schema并传入
order_schema: SchemaDict = {"price": float, "quantity": int, "product_name": str}
df = create_order_df(order_schema)

内容的提问来源于stack exchange，提问作者Joris Medeišis