Polars动态命名空间类型检查实践

本文深入探讨了在使用 Polars 动态注册 API 命名空间时，Python 类型检查器（如 Mypy 和 Pyright）报告类型错误的问题。我们将分析其根本原因，并提供两种解决方案：一是建议 Polars 官方在 `Expr` 类中添加 `__getattr__` 以实现基本抑制，二是通过构建一个 Mypy 插件来实现对动态注册命名空间的全面静态类型检查，从而在开发过程中捕获更多潜在错误。

理解 Polars 动态命名空间与类型检查器的冲突

Polars 提供了一个强大的 API，允许用户通过 @pl.api.register_expr_namespace 装饰器注册自定义表达式命名空间。这使得用户可以创建类似 pl.all().my_namespace.my_function() 这样的链式调用，极大地增强了代码的表达力和复用性。然而，这种动态注册机制对静态类型检查器构成了挑战。

当类型检查器（如 Mypy 或 Pyright）分析以下代码时：

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"]).select(
    [
        pl.all().greetings.hello(), # 类型检查器在此处报错
        pl.all().greetings.goodbye(),
    ]
))

它们会报告类似 "Expr" has no attribute "greetings" 的错误。这是因为在代码静态分析阶段，pl.Expr 对象上并没有名为 greetings 的属性，该属性是在运行时通过 Polars 的注册机制动态添加的。静态类型检查器无法预知这种运行时行为，因此会将其识别为类型错误。

解决方案一：通过 __getattr__ 提供动态属性访问提示

解决此问题的最直接方法是让 Polars 在其核心 Expr 类中为类型检查器提供一个关于动态属性访问的提示。Python 的类型系统允许通过在类中定义 __getattr__ 方法来指示存在动态属性。

具体来说，在 polars.expr.expr.Expr 类中，如果能在类型检查模式下（即 typing.TYPE_CHECKING 为 True 时）添加一个 __getattr__ 方法的定义，就可以有效地抑制类型检查器关于动态属性访问的错误。

import typing

class Expr:
    # ... 现有代码 ...
    if typing.TYPE_CHECKING:
        def __getattr__(self, attr_name: str, /) -> typing.Any: ...
    # ... 现有代码 ...

这个 __getattr__ 的定义告诉类型检查器：当尝试访问 Expr 实例上不存在的属性时，它可能会通过 __getattr__ 方法动态地返回一个 Any 类型的值。这使得类型检查器不再报错，因为它知道这个属性可能在运行时存在。

注意事项：

• 这需要 Polars 官方在库中进行修改。用户无法直接在自己的代码中为 polars.Expr 添加此方法。
• 这种方法虽然消除了 attr-defined 错误，但它提供的是 Any 类型，这意味着后续对 greetings 命名空间内方法的调用将失去静态类型检查的优势，例如，无法检查 hello() 是否接收了错误的参数。

解决方案二：针对 Mypy 的高级静态类型检查插件

对于追求更严格、更全面的静态类型检查的用户，尤其是 Mypy 用户，可以开发一个 Mypy 插件。Mypy 插件允许开发者扩展 Mypy 的行为，使其能够理解和处理特定库的复杂或动态特性。通过插件，我们可以让 Mypy 识别 Polars 的命名空间注册机制，并为注册的命名空间提供完整的静态类型信息。

一个设计良好的 Mypy 插件可以实现以下目标：

• 消除 attr-defined 错误： 像 __getattr__ 方案一样。
• 提供命名空间内部的类型检查： 能够检查 greetings.hello() 的参数是否正确，或者 greetings 命名空间下是否存在 non_existent_method()。

期望的 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")

# 假设以下代码在一个使用插件的 test.py 文件中
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"
            pl.all().asdfjkl                # Mypy 将在此处报告：`polars.expr.expr.Expr` object has no attribute `asdfjkl`
        ]
    )
)

可以看到，插件不仅解决了属性不存在的问题，还能对命名空间内部的方法调用进行详细的参数检查。

项目结构

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

project/
  mypy.ini              # Mypy 配置文件，指定插件
  mypy_polars_plugin.py # Mypy 插件实现
  test.py               # 包含 Polars 代码的测试文件

实现 Mypy 插件

1. mypy.ini 配置

在 mypy.ini 文件中，我们需要告诉 Mypy 使用我们的插件：

[mypy]
plugins = mypy_polars_plugin.py

2. mypy_polars_plugin.py 插件代码

这个文件包含了 Mypy 插件的详细实现。插件的核心在于利用 Mypy 提供的钩子（hooks）来修改其对 Polars 表达式的类型推断行为。

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

# 定义一些常量，便于引用 Polars 相关的全限定名
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 Mypy 插件实现。
    它通过 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)
        # 用于存储已注册的 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:
        """
        这个钩子用于在 MRO (Method Resolution Order) 解析时自定义类的行为。
        它被用来为 `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_registering_hook` 来记录注册的命名空间。
        """
        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` 及其子类，
        它会调用 `polars_expr_attribute_hook` 来处理动态命名空间的属性访问。
        """
        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` 装饰器的逻辑。
        它从装饰器参数中提取命名空间名称，并将其与被装饰的类的类型关联起来，
        存储在 `_polars_expr_namespace_name_to_type_dict` 中。
        """
        # 确保装饰器表达式是 `@polars.api.register_expr_namespace()`
        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__` 方法。
    这个方法仅用于类型检查，告知 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),
    )

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 test.py (确保在 project 目录下执行)，Mypy 将会按照预期报告类型错误，而不是简单的 attr-defined。

Pyright 的限制

与 Mypy 不同，Pyright 目前不提供官方的插件机制来扩展其类型检查行为。这意味着对于 Polars 动态命名空间问题，Pyright 用户只能依赖以下方法：

• 行内忽略： 在每一行报错的代码后添加 # type: ignore[attr-defined] 或 # pyright: ignore[reportGeneralTypeIssues]。
• 文件级别控制： 在文件顶部添加 # pyright: reportUnknownMemberType=none, reportGeneralTypeIssues=none 来禁用相关检查，但这会降低整个文件的类型安全性。

由于 Pyright 核心开发者对插件支持的谨慎态度，除非 Python 类型系统引入新的 PEP 来标准化动态命名空间注册的类型提示，否则 Pyright 很难提供像 Mypy 插件那样细致的静态类型检查。

总结

Polars 的动态命名空间注册功能虽然强大，但与 Python 静态类型检查器之间存在固有的冲突。解决这些冲突有多种途径：

1. Polars 官方改进： 建议 Polars 在 Expr 类中添加 typing.TYPE_CHECKING 条件下的 __getattr__ 定义，以提供基本的类型检查器兼容性，消除 attr-defined 错误。
2. Mypy 插件： 对于需要全面静态类型检查的用户，开发一个 Mypy 插件是最佳实践。通过插件，可以实现对动态注册命名空间的细致类型推断和错误报告，显著提升代码质量和可维护性。
3. Pyright 妥协： Pyright 用户目前只能通过忽略注释或文件级配置来抑制错误，无法实现像 Mypy 插件那样的深度静态分析。

在实际开发中，如果团队使用 Mypy，强烈推荐投入精力开发或寻找现有 Mypy 插件，以充分利用静态类型检查的优势。这不仅能解决当前的类型错误，还能在早期发现潜在的逻辑问题，从而提高 Polars 应用的健壮性。