如何为Polars API扩展配置MyPy以解决类型检查报错？

方法1：类型检查阶段扩展Polars类 #

借助TYPE_CHECKING常量，只在类型检查时给pl.DataFrame添加split属性的注解，完全不影响运行时逻辑：

from typing import TYPE_CHECKING
import polars as pl
from your_extension_module import SplitFrame

if TYPE_CHECKING:
    # 定义带split属性的子类
    class _DataFrameWithSplit(pl.DataFrame):
        split: SplitFrame

    # 替换类型检查时的DataFrame定义
    pl.DataFrame = _DataFrameWithSplit  # type: ignore[misc, assignment]

这段代码只会在MyPy检查时生效，实际运行时pl.DataFrame还是原来的类，不用担心中断现有功能。

方法2：用Protocol定义扩展接口 #

自定义一个Protocol，描述带有split属性的DataFrame，然后在变量或函数参数上标注这个类型：

from typing import Protocol
import polars as pl
from your_extension_module import SplitFrame

class DataFrameWithSplit(Protocol):
    split: SplitFrame

# 使用时标注类型
def process_alternate_rows(df: DataFrameWithSplit) -> tuple[pl.DataFrame, pl.DataFrame]:
    return df.split.by_alternate_rows()

# 调用时自动兼容
df = pl.read_csv("data.csv")
df1, df2 = process_alternate_rows(df)

这种方式不需要修改Polars的类，只需要在需要用到扩展属性的地方标注自定义Protocol即可，维护成本极低。

方法3：精准类型转换替代全局忽略 #

如果只是零散调用split属性，用cast做临时类型转换，比# type: ignore更严谨：

from typing import cast
import polars as pl
from your_extension_module import SplitFrame

df = pl.read_csv("data.csv")
# 转换为带split属性的类型
typed_df = cast(type(df).__class__ & DataFrameWithSplit, df)
typed_df.split.by_alternate_rows()

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