Polars动态API兼容方案解析

文章小白一枚，正在不断学习积累知识，现将学习到的知识记录一下，也是将我的所得分享给大家！而今天这篇文章《Polars动态API注册兼容性解决方案》带大家来了解一下##content_title##，希望对大家的知识积累有所帮助，从而弥补自己的不足，助力实战开发！

本文深入探讨了在使用Polars的动态API注册功能（如`@pl.api.register_expr_namespace`）时，Mypy和Pyright等类型检查器报告`attr-defined`错误的问题。文章分析了问题的根本原因，即Python静态类型系统无法识别运行时动态添加的属性。针对此问题，本文提出了Polars官方通过定义`__getattr__`来解决的理想方案，并详细介绍了Pyright的现有规避方法以及Mypy通过自定义插件实现完全静态类型检查的详细教程，包括插件结构、代码实现及效果展示，旨在帮助开发者在享受Polars灵活性的同时，维护代码的类型安全。

Polars动态API注册与类型检查器挑战

Polars提供了强大的API注册机制，允许用户为Expr等对象动态地扩展命名空间，例如通过@pl.api.register_expr_namespace装饰器。这种灵活性在运行时表现出色，但在静态类型检查阶段，Mypy或Pyright等工具会因为无法在polars.Expr类定义中找到这些动态注册的属性而报错，典型的错误是attr-defined。这是因为Python的类型系统默认是静态的，它无法预知在程序运行时才会被添加的属性。

考虑以下官方文档中的示例，它定义了一个名为greetings的表达式命名空间：

import polars as pl

@pl.api.register_expr_namespace("greetings")
class Greetings:
    def __init__(self, expr: pl.Expr):
        self._expr = expr

    def hello(self) -> pl.Expr:
        return (pl.lit("Hello ") + self._expr).alias("hi there")

    def goodbye(self) -> pl.Expr:
        return (pl.lit("Sayōnara ") + self._expr).alias("bye")

print(pl.DataFrame(data=["world", "world!", "world!!"]).select(
    [
        pl.all().greetings.hello(), # mypy/pyright会在此处报错
        pl.all().greetings.goodbye(),
    ]
))

运行Mypy或Pyright会得到如下错误：

% mypy checker.py
checker.py:19: error: "Expr" has no attribute "greetings"  [attr-defined]
Found 1 error in 1 file (checked 1 source file

% pyright checker.py
/path/to/checker.py:19:18 - error: Cannot access member "greetings" for type "Expr"
    Member "greetings" is unknown (reportGeneralTypeIssues)

这些错误表明，类型检查器无法识别pl.all()（其类型为pl.Expr）上动态注册的greetings属性。

Polars层面的理想解决方案：引入 __getattr__

解决此问题的最根本且理想的方式是Polars库自身在polars.expr.expr.Expr类中定义一个特殊的__getattr__方法，并结合typing.TYPE_CHECKING标志。__getattr__是一个钩子，当访问一个不存在的属性时会被调用。类型检查器可以利用它的存在来推断动态属性访问的可能性。

在Expr类中添加类似以下结构的代码，足以让类型检查器停止对动态属性访问的报错：

import typing

class Expr:
    # ... Expr类的其他定义 ...

    if typing.TYPE_CHECKING:
        def __getattr__(self, attr_name: str, /) -> typing.Any: ...

这个if typing.TYPE_CHECKING:块确保了__getattr__只在类型检查时可见，不会影响运行时行为。它向类型检查器发出信号：Expr对象可能会在运行时拥有任何属性，并且这些属性的类型是Any。这虽然不能提供具体的类型信息，但能有效消除attr-defined错误。

建议Polars开发者考虑在库中添加此类声明，以提升与类型检查工具的兼容性。

Pyright的局限性与应对策略

Pyright作为一个强大的类型检查器，其设计哲学决定了它对插件机制持谨慎态度。这意味着目前Pyright不支持像Mypy那样通过自定义插件来理解Polars的动态命名空间注册。

因此，对于Pyright用户，主要的规避策略包括：

1. 行内忽略注释：在每一行引发错误的表达式后添加# type: ignore[attr-defined]或# pyright: ignore[reportGeneralTypeIssues]。这虽然有效，但会使得代码中充斥着忽略注释，降低可读性。
2. 文件级别类型控制：在文件顶部添加控制注释，例如# pyright: reportUnknownMemberType=none, reportGeneralTypeIssues=none。这种方法会禁用特定类型的诊断报告，但其副作用是可能隐藏文件中其他真正的类型问题。

这些方法都是临时的权宜之计，无法提供真正的静态类型安全。

Mypy的静态类型检查方案：自定义插件

相比Pyright，Mypy提供了强大的插件系统，允许开发者扩展其类型推断能力。通过编写一个Mypy插件，我们可以让Mypy“理解”Polars的动态命名空间注册，从而实现对自定义命名空间的完全静态类型检查。这意味着Mypy不仅不会报错，还能检查自定义命名空间内方法的参数数量、类型等。

期望的Mypy静态类型检查结果

通过Mypy插件，我们可以实现以下精确的类型检查效果：

import polars as pl

@pl.api.register_expr_namespace("greetings")
class Greetings:
    def __init__(self, expr: pl.Expr):
        self._expr = expr

    def hello(self) -> pl.Expr:
        return (pl.lit("Hello ") + self._expr).alias("hi there")

    def goodbye(self) -> pl.Expr:
        return (pl.lit("Sayōnara ") + self._expr).alias("bye")

print(
    pl.DataFrame(data=["world", "world!", "world!!"]).select(
        [
            pl.all().greetings.hello(),
            pl.all().greetings.goodbye(1),  # Mypy: Too many arguments for "goodbye" of "Greetings"  [call-arg]
            pl.all().asdfjkl                # Mypy: `polars.expr.expr.Expr` object has no attribute `asdfjkl`  [misc]
        ]
    )
)

如上所示，插件不仅能识别greetings命名空间，还能正确地指出goodbye(1)的参数错误以及asdfjkl这个不存在的属性。

项目结构

为了实现Mypy插件，我们需要一个特定的项目结构：

project/
  mypy.ini
  mypy_polars_plugin.py
  test.py

插件实现详解

1. mypy.ini 配置

在mypy.ini文件中，我们需要告诉Mypy加载我们的自定义插件：

[mypy]
plugins = mypy_polars_plugin.py

2. mypy_polars_plugin.py 插件代码

这是插件的核心，它通过Mypy提供的钩子来扩展类型检查逻辑。

from __future__ import annotations

import typing_extensions as t

import mypy.nodes
import mypy.plugin
import mypy.plugins.common

if t.TYPE_CHECKING:
    import collections.abc as cx

    import mypy.options
    import mypy.types

STR___GETATTR___NAME: t.Final = "__getattr__"
STR_POLARS_EXPR_MODULE_NAME: t.Final = "polars.expr.expr"
STR_POLARS_EXPR_FULLNAME: t.Final = f"{STR_POLARS_EXPR_MODULE_NAME}.Expr"
STR_POLARS_EXPR_REGISTER_EXPR_NAMESPACE_FULLNAME: t.Final = "polars.api.register_expr_namespace"

def plugin(version: str) -> type[PolarsPlugin]:
    """Mypy插件的入口点。"""
    return PolarsPlugin

class PolarsPlugin(mypy.plugin.Plugin):

    _polars_expr_namespace_name_to_type_dict: dict[str, mypy.types.Type]

    def __init__(self, options: mypy.options.Options) -> None:
        super().__init__(options)
        # 用于存储已注册的Polars表达式命名空间名称及其对应的类型
        self._polars_expr_namespace_name_to_type_dict = {}

    @t.override
    def get_customize_class_mro_hook(
        self, fullname: str
    ) -> cx.Callable[[mypy.plugin.ClassDefContext], None] | None:
        """
        这个钩子允许在Mypy处理类的MRO（方法解析顺序）之前修改类定义。
        我们利用它为`polars.expr.expr.Expr`类动态添加一个虚拟的`__getattr__`方法，
        以满足Mypy在`get_attribute_hook`工作时对动态属性访问的最低要求。
        """
        if fullname == STR_POLARS_EXPR_FULLNAME:
            return add_getattr
        return None

    @t.override
    def get_class_decorator_hook_2(
        self, fullname: str
    ) -> cx.Callable[[mypy.plugin.ClassDefContext], bool] | None:
        """
        此钩子在Mypy遇到类装饰器时触发。
        我们关注`@polars.api.register_expr_namespace(...)`装饰器，
        从中提取命名空间名称，并将其与被装饰的类（即命名空间类）的类型关联起来，
        存储在`_polars_expr_namespace_name_to_type_dict`中。
        """
        if fullname == STR_POLARS_EXPR_REGISTER_EXPR_NAMESPACE_FULLNAME:
            return self.polars_expr_namespace_registering_hook
        return None

    @t.override
    def get_attribute_hook(
        self, fullname: str
    ) -> cx.Callable[[mypy.plugin.AttributeContext], mypy.types.Type] | None:
        """
        当Mypy需要解析一个属性访问（例如`expr.greetings`）时，此钩子被调用。
        如果被访问的对象是`polars.expr.expr.Expr`的实例，并且该属性在
        `_polars_expr_namespace_name_to_type_dict`中注册过，
        我们就返回对应命名空间类的类型，从而实现静态类型检查。
        """
        if fullname.startswith(f"{STR_POLARS_EXPR_FULLNAME}."):
            return self.polars_expr_attribute_hook
        return None

    def polars_expr_namespace_registering_hook(
        self, ctx: mypy.plugin.ClassDefContext
    ) -> bool:
        """
        处理`@polars.api.register_expr_namespace`装饰器。
        从装饰器参数中解析出命名空间名称，并将命名空间类（`ctx.cls`）的类型存储起来。
        """
        namespace_arg: str | None
        if (
            (not isinstance(ctx.reason, mypy.nodes.CallExpr))
            or (len(ctx.reason.args) != 1)
            or (
                (namespace_arg := ctx.api.parse_str_literal(ctx.reason.args[0])) is None
            )
        ):
            # 如果装饰器表达式不符合预期（例如参数不是单个字符串字面量），则提前返回。
            return True

        self._polars_expr_namespace_name_to_type_dict[
            namespace_arg
        ] = ctx.api.named_type(ctx.cls.fullname)

        return True

    def polars_expr_attribute_hook(
        self, ctx: mypy.plugin.AttributeContext
    ) -> mypy.types.Type:
        """
        处理`polars.expr.expr.Expr`实例上的属性访问。
        如果属性名对应一个已注册的命名空间，则返回该命名空间类的类型；
        否则，Mypy会报告一个错误，指示`Expr`对象没有该属性。
        """
        assert isinstance(ctx.context, mypy.nodes.MemberExpr)
        attr_name: str = ctx.context.name
        namespace_type: mypy.types.Type | None = (
            self._polars_expr_namespace_name_to_type_dict.get(attr_name)
        )
        if namespace_type is not None:
            return namespace_type
        else:
            ctx.api.fail(
                f"`{STR_POLARS_EXPR_FULLNAME}` object has no attribute `{attr_name}`",
                ctx.context,
            )
            return mypy.types.AnyType(mypy.types.TypeOfAny.from_error)


def add_getattr(ctx: mypy.plugin.ClassDefContext) -> None:
    """
    辅助函数，用于向指定的类（这里是`Expr`）添加一个虚拟的`__getattr__`方法。
    """
    mypy.plugins.common.add_method_to_class(
        ctx.api,
        cls=ctx.cls,
        name=STR___GETATTR___NAME,
        args=[
            mypy.nodes.Argument(
                variable=mypy.nodes.Var(
                    name="attr_name", type=ctx.api.named_type("builtins.str")
                ),
                type_annotation=ctx.api.named_type("builtins.str"),
                initializer=None,
                kind=mypy.nodes.ArgKind.ARG_POS,
                pos_only=True,
            )
        ],
        return_type=mypy.types.AnyType(mypy.types.TypeOfAny.implementation_artifact),
        self_type=ctx.api.named_type(STR_POLARS_EXPR_FULLNAME),
    )

3. test.py 示例代码

此文件与之前示例相同，用于验证Mypy插件的效果：

import polars as pl


@pl.api.register_expr_namespace("greetings")
class Greetings:
    def __init__(self, expr: pl.Expr):
        self._expr = expr

    def hello(self) -> pl.Expr:
        return (pl.lit("Hello ") + self._expr).alias("hi there")

    def goodbye(self) -> pl.Expr:
        return (pl.lit("Sayōnara ") + self._expr).alias("bye")


print(
    pl.DataFrame(data=["world", "world!", "world!!"]).select(
        [
            pl.all().greetings.hello(),
            pl.all().greetings.goodbye(1),  # Mypy现在会报错
            pl.all().asdfjkl                # Mypy现在会报错
        ]
    )
)

通过上述Mypy插件，开发者可以为Polars的动态API注册功能获得全面的静态类型检查支持，极大地提升了代码的健壮性和可维护性。

总结与建议

Polars的动态API注册机制为数据操作提供了极大的灵活性，但其与Python静态类型检查器的兼容性问题是开发者面临的常见挑战。

1. Polars官方改进：最理想的解决方案是Polars库在Expr等核心类中通过if typing.TYPE_CHECKING: def __getattr__来明确告知类型检查器动态属性的存在。建议社区向Polars开发者提出此功能请求。
2. Pyright用户：由于Pyright不支持插件，目前只能依赖于行内忽略注释或文件级别诊断控制来规避错误，但这会牺牲部分类型安全。
3. Mypy用户：可以利用Mypy强大的插件系统，如本文所示，实现对Polars动态命名空间的全面静态类型检查。这不仅能消除attr-defined错误，还能对动态注册的方法进行参数和返回类型检查，提供最高级别的类型安全保障。

选择哪种方案取决于项目的具体需求、使用的类型检查器以及对类型安全的要求。对于追求极致类型安全和良好开发体验的Mypy用户，自定义插件无疑是最佳选择。

文中关于的知识介绍，希望对你的学习有所帮助！若是受益匪浅，那就动动鼠标收藏这篇《Polars动态API兼容方案解析》文章吧，也可关注golang学习网公众号了解相关技术文章。