本文深入探讨了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目前不支持插件系统来处理这种动态类型注册。这意味着除了以下两种方式外，没有其他方法可以抑制这些错误：
- 行内忽略： 在每一行报错的代码后添加# type: ignore[attr-defined]或# pyright: ignore[reportGeneralTypeIssues]。
- 文件级控制： 在文件顶部添加指令，如# 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 插件实现
这个插件的核心思想是：
- 在类型检查阶段，为polars.expr.expr.Expr类动态添加一个__getattr__方法，以允许Mypy进行属性查找。
- 识别@pl.api.register_expr_namespace装饰器，并记录注册的命名空间名称及其对应的类类型。
- 当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、早期发现类型错误）是显著的。
