输入不可输入的内容:生成 Python .pyi 存根

在Python中,一些类动态地定义它们的属性。这意味着这些属性并不存在于正常的实例变量中,并且通常对IDE和静态类型检查器(如mypy)是不可见的。

一个抽象的例子:

class BookObject:
    fields = ("author", "title")

    def __init__(self, data: dict):
        self._data = {}
        for field in self.fields:

self._data[field] = data.get(field)

def __getattr__(self, name):
    if name in self._data:
        return self._data[name]
    raise AttributeError

你可以这样使用它:

book = BookObject({"title": "沙丘", "author": "弗兰克·赫伯特"})
print(book.title)

 

它在运行时有效,但你的IDE无法自动补全 book.title,而且 mypy 也无法检查它的类型。

我在 oslo.versionedobjects 中看到了这个问题——它在OpenStack中用于数据模型。

这个库允许你这样定义字段:

class Book(BaseObject):
    fields = {
        "id": fields.IntegerField(),
"标题": fields.StringField(),
"作者": fields.StringField(),
"评分": fields.IntegerField(),
}

在运行时,您可以像访问普通属性一样访问这些字段(但在背后发生了很多魔法),然而 IDE 和 mypy 无法看到它们。

然后我有了一个主意:为了好玩,我可以为项目添加自动 .pyi 存根文件生成,这样在 IDE 中工作会更容易,并且 mypy 可以找到类型错误。

.pyi 文件

存根文件(.pyi)在 PEP 484 中进行了描述。
它们仅包含类型提示,静态分析器使用它们。

示例:

class Book:
    title: str
    author: str

即使真实的类属性是动态的,存根文件也能让工具知道哪些属性存在及其类型。

演示项目

我制作了一个小型演示项目。完整代码在这里:演示仓库

项目结构:

├── main.py
├── Makefile
├── objects
│ ├── __init__.py
│ ├── base.py
│ └── book.py
├── pyproject.toml
└── stubgen.py

Book 类看起来是这样的:

@DemoObjectRegistry.register
class Book(BaseVersionedObject):
    VERSION = "1.0"

    fields = {
        "id": fields.IntegerField(),
        "title": fields.StringField(),
{
    "author": fields.StringField(),
    "rating": fields.IntegerField(),
}
main.py 中有一个使用示例:
def echo(title: str, author_name: str, rating: str) -> None:
print(f"{title.upper()} by {author_name.upper()} with rating {rating.strip()}!")

def main() -> None:
books = [

以下是您提供的英文技术文章内容的中文翻译:

书籍(id=1, 标题=注定的城市, 作者=阿尔卡季与鲍里斯·斯特鲁加茨基, 评分=10),
书籍(id=2, 标题=定居点, 作者=基尔·布柳切夫, 评分=10),
]

对于 b 书籍:

请注意,翻译保持了技术术语的准确性和一致性,同时保留了原有的HTML格式和代码块。

echo(b.title, b.author, b.rating)

注意:echo期望ratingstr类型,但Book.ratingint类型。

没有存根,mypy无法检测到这一点:

$ make lint
uv run ruff check .
所有检查通过!
uv run mypy .

成功:在 5 个源文件中未发现问题  
未检测到任何内容,因为动态字段是不可见的。

生成 .pyi 存根

为了解决这个问题,我编写了一个小的 Python 脚本 stubgen.py,它自动生成带有适当类型提示的 .pyi 存根。

以下是它的工作原理。

步骤 1 - 生成基线存根

首先,我们使用 mypy stubgen 生成初始存根。

这会生成包含类和方法的 .pyi 文件,但动态字段缺失:

class Book(BaseVersionedObject):
    VERSION: str
    fields: Incomplete
步骤 2 - 使用 ast 解析源代码并映射字段类型

然后,脚本使用 Python 的 ast 模块解析原始 Python 源代码,并在每个类中查找字段字典。

对于每个字段,脚本确定其 Python 类型(intstr 等)。

每个字段构造函数(如 IntegerFieldStringField)都映射到一个 Python 类型:

FIELD_TYPE_MAP = {
    "IntegerField": "int",
    "StringField": "str",
    "BooleanField": "bool",
    "FloatField": "float",
}
步骤3 - 更新 .pyi 存根

最后,脚本通过使用 ast 插入类型属性来更新 .pyi 存根:

# 更新前
class Book(BaseVersionedObject):
    VERSION: str
    fields: Incomplete

# 更新后
class Book(BaseVersionedObject): rating: int author: str title: str id: int VERSION: str fields: Incomplete

运行 stubgen 和 linters

现在我们可以通过运行以下命令生成 .pyi 存根:

$ make stubgen

接下来,我们运行 linters:

$ make lint
uv run ruff check .
所有检查通过!
uv run mypy .
main.py:15: error: Argument 3 to "echo" has incompatible type "int"; expected "str"  [arg-type]

发现 1 个错误 1 个文件中 (检查了 5 个源 文件)
make: *** [lint] 错误 1

多亏了生成的存根,mypy 现在检测到 Book.rating 是一个 int,而函数 echo 期望的是一个字符串。

结论

这当然是一个简化的例子 - 该脚本没有处理更复杂的情况,如 OptionalDictList 或嵌套对象。

但主要思想是,您可以在自己的项目中使用这种方法,使处理动态定义的属性变得更加容易。

更多