Polars命名空间优化方案解析

本文深入探讨了Polars库自定义命名空间（`@pl.api.register_expr_namespace`）与Python静态类型检查器（如Mypy和Pyright）的兼容性难题。Polars的动态属性注册机制导致类型检查器常报`attr-defined`错误，影响代码质量和开发效率。文章提出两种优化方案：一是建议Polars在`Expr`类中添加类型检查专用的`__getattr__`定义，直接从库层面解决；二是针对Mypy，提供详细的插件实现，实现自定义命名空间的完整静态类型检查，消除手动`# type: ignore`的需要，提升代码健壮性。尽管Mypy插件实现相对复杂，但长期收益显著，能有效减少类型错误，提升开发效率。文章旨在为Polars用户提供清晰的解决方案，助力构建更稳定、更易维护的Polars应用。

本文深入探讨了Polars库中自定义命名空间（`@pl.api.register_expr_namespace`）与Python静态类型检查器（如Mypy和Pyright）之间的兼容性问题。由于Polars的动态属性注册机制，类型检查器通常会报告`attr-defined`错误。文章提出了两种主要解决方案：一是建议Polars在`Expr`类中添加类型检查专用的`__getattr__`定义；二是为Mypy提供一个详细的插件实现，以实现自定义命名空间的完整静态类型检查，从而消除手动`# type: ignore`的需要。

Polars自定义命名空间与类型检查挑战

Polars是一个高性能的DataFrame库，它允许用户通过@pl.api.register_expr_namespace装饰器注册自定义表达式命名空间，极大地扩展了其功能。然而，这种动态的属性注册方式给静态类型检查带来了挑战。当您尝试通过自定义命名空间访问方法时，例如pl.all().greetings.hello()，Mypy或Pyright等类型检查器会因为无法在polars.Expr类型上静态地找到greetings属性而报告attr-defined错误。

以下是一个典型的示例，展示了使用自定义命名空间时类型检查器报告的错误：

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 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)
1 error, 0 warnings, 0 informations

这些错误源于Polars的动态注册机制与Python静态类型系统之间的不匹配。类型检查器在编译时无法预知Expr对象上将存在哪些动态注册的属性。

解决方案一：Polars库层面的改进

最直接且理想的解决方案是由Polars库本身在polars.expr.expr.Expr类中添加一个类型检查专用的__getattr__定义。Python的typing模块提供了TYPE_CHECKING常量，允许在类型检查阶段引入特定的类型提示，而不会影响运行时行为。

通过在Expr类中添加如下结构：

import typing

class Expr:
    # ... 现有代码 ...

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

这个__getattr__的定义会向类型检查器发出信号，表明Expr类支持动态属性访问，并且这些属性的类型可以是Any。这足以抑制关于动态属性访问的attr-defined错误，而无需在代码中手动添加# type: ignore。我们建议向Polars开发者提出此功能请求，以改善其类型提示的兼容性。

解决方案二：针对不同类型检查器的处理

在Polars尚未实现上述改进的情况下，我们可以根据所使用的类型检查器采取不同的策略。

Pyright的限制与临时对策

Pyright目前不支持插件系统来处理这种动态类型注册。这意味着除了以下两种方式外，没有其他方法可以抑制这些错误：

1. 行内忽略： 在每一行报错的代码后添加# type: ignore[attr-defined]或# pyright: ignore[reportGeneralTypeIssues]。
2. 文件级控制： 在文件顶部添加指令，如# pyright: reportUnknownMemberType=none, reportGeneralTypeIssues=none，但这会禁用对整个文件中相关类型问题的报告，可能掩盖其他潜在问题。

由于Pyright的架构设计，其维护者对插件支持持谨慎态度，因此在短期内，Pyright用户可能需要依赖这些手动忽略指令。

Mypy插件实现完整静态类型检查

Mypy拥有强大的插件系统，这使得我们可以为Polars的动态命名空间注册提供一个静态类型检查的解决方案。通过编写一个Mypy插件，我们可以在类型检查阶段“教导”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属性的attr-defined错误，还能正确地检查goodbye方法的参数错误和不存在的asdfjkl属性。

项目结构

为了实现Mypy插件，我们需要以下文件结构：

project/
  mypy.ini
  mypy_polars_plugin.py
  test.py

mypy.ini 配置

在 mypy.ini 文件中，我们需要指定Mypy加载我们的插件：

[mypy]
plugins = mypy_polars_plugin.py

mypy_polars_plugin.py 插件实现

这个插件的核心思想是：

1. 在类型检查阶段，为polars.expr.expr.Expr类动态添加一个__getattr__方法，以允许Mypy进行属性查找。
2. 识别@pl.api.register_expr_namespace装饰器，并记录注册的命名空间名称及其对应的类类型。
3. 当Mypy尝试访问polars.expr.expr.Expr实例上的属性时，如果该属性是已注册的命名空间，则返回其对应的类类型。

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):
    """
    一个Mypy插件，用于处理Polars自定义表达式命名空间的类型检查。
    """

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

    def __init__(self, options: mypy.options.Options) -> None:
        super().__init__(options)
        # 用于存储已注册的命名空间名称到其类型对象的映射
        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在`Expr`实例上查找动态属性。
        """
        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(...)`装饰器，
        并记录被装饰的类（即自定义命名空间类）的类型。
        """
        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解析类实例的属性访问时被调用。
        它用于处理`polars.expr.expr.Expr`实例上的属性访问，
        如果属性名对应一个已注册的命名空间，则返回该命名空间的类型。
        """
        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(<namespace name>)`装饰器。
        它从装饰器参数中提取命名空间名称，并将其与被装饰的类类型关联起来，
        存储在`_polars_expr_namespace_name_to_type_dict`中。
        """
        # 确保装饰器表达式是`@polars.api.register_expr_namespace(<namespace name>)`的形式
        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:
        """
        当Mypy访问`polars.expr.expr.Expr`实例上的属性时，此方法会被调用。
        它会检查被访问的属性名是否在已注册的命名空间字典中。
        如果是，则返回该命名空间的类型；否则，Mypy会报告一个错误。
        """
        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:
    """
    一个辅助函数，用于向给定的类定义（`polars.expr.expr.Expr`）添加一个虚拟的`__getattr__`方法。
    这个`__getattr__`方法接受一个字符串`attr_name`并返回`Any`类型，
    从而告诉Mypy该类支持动态属性访问。
    """
    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),
    )

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: Too many arguments for "goodbye" of "Greetings"  [call-arg]
            pl.all().asdfjkl                # Mypy: `polars.expr.expr.Expr` object has no attribute `asdfjkl`
        ]
    )
)

通过这种Mypy插件的实现，开发者可以获得Polars自定义命名空间的完整静态类型检查能力，极大地提升了代码质量和开发效率。

总结

Polars的动态API注册机制为扩展功能提供了便利，但也带来了与静态类型检查器兼容性的挑战。理想情况下，Polars库应在其Expr类中添加一个类型检查专用的__getattr__定义来解决这一问题。在此之前，Pyright用户可能需要依赖手动忽略指令。对于Mypy用户，通过实现一个自定义插件，可以完全解决类型检查问题，实现对自定义命名空间的精确静态分析，从而避免运行时错误并提升代码的健壮性。虽然Mypy插件的实现相对复杂，但其带来的长期收益（减少# type: ignore、早期发现类型错误）是显著的。

本篇关于《Polars命名空间优化方案解析》的介绍就到此结束啦，但是学无止境，想要了解学习更多关于文章的相关知识，请关注golang学习网公众号！