python 是否可以通过编程从示例化的类生成一个pyi文件?

dwbf0jvd  于 2023-02-11  发布在  Python
关注(0)|答案(1)|浏览(117)

我从字典中创建了一个类,如下所示:

class MyClass:
    def __init__(self, dictionary):
        for k, v in dictionary.items():
            setattr(self, k, v)

我正在尝试弄清楚如何为这个动态生成的类获取智能感知,大多数IDE都可以读取pyi文件来做这类事情。
不过我不想手工写出一个pyi文件。
是否可以示例化这个类并通过编程将一个pyi文件从它写入磁盘?
mypy有stubgen工具,但我不知道是否可以这样使用它。
我可以从mypy导入stubgen并以某种方式给它提供MyClass(<some dict>)吗?

lskq00tm

lskq00tm1#

像stubgen这样的静态分析程序是分析动态填充类的错误工具,因为它们无法看到完整类的源代码来给予类的存根,必须在运行时通过运行源代码来首先填充示例属性来生成存根。
假设您有一个动态填充的类,如您的示例所示,

class MyClass:
    def __init__(self, dictionary: dict[str, object]) -> None:
        k: str
        v: object
        for k, v in dictionary.items():
            setattr(self, k, v)

把这个字典传递给构造函数

import statistics

instance: MyClass = MyClass({"a": 1, "b": "my_string", "distribution": statistics.NormalDist(0.0, 1.0)})

你想把这个作为你的输出

import statistics

class MyClass:
    a: int
    b: str
    distribution: statistics.NormalDist

    def __init__(self, dictionary: dict[str, object]) -> None:
        ...

生成上述输出的最简单方法是挂钩示例创建和初始化,这样就不会影响类中已经存在的__new____init__链接的super调用,这可以通过元类的__call__方法来实现:

class _PostInitialisationMeta(type):

    """
    Metaclass for classes subject to dynamic stub generation
    """

    def __call__(
        cls, dictionary: dict[str, object], *args: object, **kwargs: object
    ) -> object:

        """
        Override instance creation and initialisation. Generate a string representing
        the class's stub definition suitable for a `.pyi` file.

        Parameters
        ----------
        dictionary
            Mapping from instance attribute names to attribute values
        *args
        **kwargs
            Other positional and keyword arguments to the class's `__new__` and
            `__init__` methods

        Returns
        -------
        object
            Created instance
        """

        instance: object = super().__call__(dictionary, *args, **kwargs)
        <generate string here>
        return instance

然后,你可以将这个类解析为abstract syntax tree,通过添加、移除或转换节点来修改树,然后对转换后的树进行反解析,下面是使用Python标准库ast.NodeVisitor的一个可能的实现:

仅限Python 3.9及以上版本

from __future__ import annotations

import ast
import inspect
import typing as t

if t.TYPE_CHECKING:

    class _SupportsBodyStatements(t.Protocol):
        body: list[ast.stmt]

_CLASS_TO_STUB_SOURCE_DICT: t.Final[dict[type, str]] = {}

class _PostInitialisationMeta(type):

    """
    Metaclass for classes subject to dynamic stub generation
    """

    def __call__(
        cls, dictionary: dict[str, object], *args: object, **kwargs: object
    ) -> object:

        """
        Override instance creation and initialisation. The first time an instance of a
        class is created and initialised, cache a string representing the class's stub
        definition suitable for a `.pyi` file.

        Parameters
        ----------
        dictionary
            Mapping from instance attribute names to attribute values
        *args
        **kwargs
            Other positional and keyword arguments to the class's `__new__` and
            `__init__` methods

        Returns
        -------
        object
            Created instance
        """

        instance: object = super().__call__(dictionary, *args, **kwargs)
        _DynamicClassStubsGenerator.cache_stub_for_dynamic_class(cls, dictionary)
        return instance

def _remove_docstring(node: _SupportsBodyStatements, /) -> None:

    """
    Removes a docstring node if it exists in the given node's body
    """

    first_node: ast.stmt = node.body[0]
    if (
        isinstance(first_node, ast.Expr)
        and isinstance(first_node.value, ast.Constant)
        and (type(first_node.value.value) is str)
    ):
        node.body.pop(0)

def _replace_body_with_ellipsis(node: _SupportsBodyStatements, /) -> None:

    """
    Replaces the body of a given node with a single `...`
    """

    node.body[:] = [ast.Expr(ast.Constant(value=...))]

class _DynamicClassStubsGenerator(ast.NodeVisitor):

    """
    Generate and cache stubs for class instances whose instance variables are populated
    dynamically
    """

    @classmethod
    def cache_stub_for_dynamic_class(
        StubsGenerator, Class: type, dictionary: dict[str, object], /
    ) -> None:

        # Disallow stubs generation if the stub source is already generated
        try:
            _CLASS_TO_STUB_SOURCE_DICT[Class]
        except KeyError:
            pass
        else:
            return

        # Get class's source code
        src: str = inspect.getsource(Class)
        module_tree: ast.Module = ast.parse(src)

        class_statement: ast.stmt = module_tree.body[0]
        assert isinstance(class_statement, ast.ClassDef)

        # Strip unnecessary details from class body
        stubs_generator: _DynamicClassStubsGenerator = StubsGenerator()
        stubs_generator.visit(module_tree)

        # Adds the following:
        #  - annotated instance attributes on the class body
        #  - import statements for non-builtins
        # --------------------------------------------------
        added_import_nodes: list[ast.stmt] = []
        added_class_nodes: list[ast.stmt] = []
        k: str
        v: object
        for k, v in dictionary.items():
            value_type: type = type(v)
            value_type_name: str = value_type.__qualname__
            value_type_module_name: str = value_type.__module__

            annotated_assignment_statement: ast.stmt = ast.parse(
                f"{k}: {value_type_name}"
            ).body[0]
            assert isinstance(annotated_assignment_statement, ast.AnnAssign)
            added_class_nodes.append(annotated_assignment_statement)
            if value_type_module_name != "builtins":
                annotation_expression: ast.expr = (
                    annotated_assignment_statement.annotation
                )
                assert isinstance(annotation_expression, ast.Name)
                annotation_expression.id = (
                    f"{value_type_module_name}.{annotation_expression.id}"
                )
                added_import_nodes.append(
                    ast.Import(names=[ast.alias(name=value_type_module_name)])
                )

        module_tree.body[:] = [*added_import_nodes, *module_tree.body]
        class_statement.body[:] = [*added_class_nodes, *class_statement.body]
        _CLASS_TO_STUB_SOURCE_DICT[Class] = ast.unparse(module_tree)

    def visit_ClassDef(self, node: ast.ClassDef) -> None:
        _remove_docstring(node)
        node.keywords = []  # Clear metaclass and other keywords in class definition
        self.generic_visit(node)

    def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
        _replace_body_with_ellipsis(node)

    def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> None:
        _replace_body_with_ellipsis(node)

然后,您可以像往常一样运行类,然后检查该高速缓存_CLASS_TO_STUB_SOURCE_DICT中存储的内容:

>>> MyClass({"a": 1, "b": "my_string", "distribution": statistics.NormalDist(0.0, 1.0)})
>>> src: str
>>> for src in _CLASS_TO_STUB_SOURCE_DICT.values():
...     print(src)
...
import statistics

class MyClass:
    a: int
    b: str
    distribution: statistics.NormalDist

    def __init__(self, dictionary: dict[str, object]) -> None:
        ...

在实践中,.pyi文件在 per-module 的基础上形成类型接口,因此上面的实现不能立即使用,因为它只适用于一个类。您还必须对.pyi模块中的其他类型的节点做更多的处理,决定如何处理未注解的节点、重复导入等。在将源代码写入.pyi文件之前,您可以使用stubgen,它可以分析模块的静态部分,然后您可以获取该输出并编写ast.NodeTransformer,将该输出转换为动态生成的类。

相关问题