如何为 Polars API 扩展配置 MyPy？

Question

我正在围绕

polars

Series 和 DataFrame 编写数据操作代码。 为了让我的生活更轻松，代码更具可读性，我使用 Polars API 扩展:

@pl.api.register_dataframe_namespace("split")
class SplitFrame:
    def __init__(self, df: pl.DataFrame):
        self._df = df

    def by_alternate_rows(self) -> list[pl.DataFrame]:
        df = self._df.with_row_count(name="n")
        return [
            df.filter((pl.col("n") % 2) == 0).drop("n"),
            df.filter((pl.col("n") % 2) != 0).drop("n"),
        ]


pl.DataFrame(
    data=["aaa", "bbb", "ccc", "ddd", "eee", "fff"],
    columns=[("txt", pl.Utf8)],
).split.by_alternate_rows()

# [┌─────┐  ┌─────┐
#  │ txt │  │ txt │
#  │ --- │  │ --- │
#  │ str │  │ str │
#  ╞═════╡  ╞═════╡
#  │ aaa │  │ bbb │
#  │ ccc │  │ ddd │
#  │ eee │  │ fff │
#  └─────┘, └─────┘]

我的问题是我的代码库是由 MyPy 检查的，它一点也不喜欢它们：

error: "DataFrame" has no attribute "split"  [attr-defined]

目前我能想到两种解决方法，但都不满意：

1. # type: ignore

   评论
   这就是我目前正在做的事情，但它们需要无处不在并破坏一点代码库。
2. 存根文件
   我没有探索这一途径，但我知道可以为库提供存根文件。我的问题是，Polar 和我的扩展程序都在积极开发，因此我不想“冻结”它们的签名，因为我认为维护这些存根将是一个真正的负担。

本质上我想告诉 MyPy “请假设 Polars.DataFrame 有一个

split

类型的

SplitFrame

属性。” 就是这样。 任何关于如何做到这一点或具有相同目的的东西的建议将不胜感激。

Answer 1

存根文件通常是这里的方法，尽管保持最新可能有点麻烦。我们的代码库中有一个重要的动态生成语法，并通过使用

mypy

Stubgen 的预提交挂钩来解决它。首先，它生成通常的静态类型，但然后我们导入/内省相关类的运行时版本，并将动态注册的方法注入到存根中。做到这一点确实需要一些工作，但是一旦完成，您就再也不用碰它了，所以如果您积极开发/扩展代码，那么这些努力肯定会得到回报。

这是为模块生成标准存根的示例；然后您可以读取该文件并将动态分配的方法注入其中。

from mypy.stubgen import Options as StubOptions, generate_stubs

module_file = os.path.abspath( module_path )
module_stubfile = f"{module_file}i"

# create the standard _static_ typing stub
generate_stubs(
    StubOptions(
        # 'module_file' contains the base class 
        #  onto which dynamic methods are registered
        files = [module_file],
        include_private = True,
        export_less = False,

        # standard params (not auto-defaulted)
        pyversion = sys.version_info[:2],
        interpreter = sys.executable,
        output_dir = '',
        ignore_errors = False,
        parse_only = True,
        no_import = True,
        search_path = [],
        doc_dir = '',
        packages = [],
        modules = [],
        verbose = False,
        quiet = False,
    )
)

也可能值得研究

mypy

插件挂钩，尽管老实说我无法说出它们在这种情况下的有效性：https://mypy.readthedocs.io/en/stable/extending_mypy.html

Answer 2

解决方法是全局忽略

attr-defined

错误代码。您可以通过将以下内容添加到您的

pyproject.toml

中来实现此目的：

[tool.mypy]
disable_error_code = ["attr-defined"]

另一个选项是完全禁用 Polars 库的类型检查：

[[tool.mypy.overrides]]
module = "polars"
follow_imports = "skip"

这些选项都感觉不好，但它们会抑制您看到的误报。