测试

Question

我需要提取在mypy类型的文件中导入的模块的名称，例如：

from typing import TYPE_CHECKING
if TYPE_CHECKING:
    import abc
    from django.utils import timezone

基于上面的示例，“名称提取器”功能应返回django.utils.timezone和abc，但不应返回typing.TYPE_CHECKING。

我能想到的唯一方法和超级hacky神奇的方法是使用ast和compile库：

import ast

code = """
from typing import TYPE_CHECKING
if TYPE_CHECKING:
    import abc
    from django.utils import timezone

if 'aaa':
    print('bbb')

print('hello world')
"""
tree = ast.parse(code)

# removing all nodes except for "if TYPE_CHECKING"
tree.body = [
    b for b in tree.body
    if isinstance(b, ast.If)
       and isinstance(b.test, ast.Name)
       and b.test.id == 'TYPE_CHECKING'
]

compiled = compile(tree, filename="<ast>", mode="exec")
print(compiled.co_names)

有适当的方法吗？

Answer 1

您最初使用ast的方法是一个好的开始，但是您可能最终会错过一些极端的情况（例如我在评论中已经提到的情况），这可以做很多工作避免。

如果您不介意导入模块，而您正在分析（并且我们确定存在django之类的所有子模块），那么我们可以采用“肮脏”的方式进行此操作：

将TYPE_CHECKING设置为False（默认）的导入模块。
将TYPE_CHECKING设置为True的导入模块。
比较2个模块并获取“额外”子模块名称。

首先让我们定义用于猴子修补的实用程序功能

from contextlib import contextmanager


@contextmanager
def patch(object_, attribute_name, value):
    old_value = getattr(object_, attribute_name)
    try:
        setattr(object_, attribute_name, value)
        yield
    finally:
        setattr(object_, attribute_name, old_value)

之后，可以使用类似{p>的stdlib中的importlib模块（用于动态导入和重新加载）和inspect模块（用于检查对象是否为模块）来编写我们的函数。

import importlib
import inspect
import typing


def detect_type_checking_mode_modules_names(module_name):
    module = importlib.import_module(module_name)
    default_module_names = set(vars(module))

    with patch(typing, 'TYPE_CHECKING', True):
        # reloading since ``importlib.import_module``
        # will return previously cached entry
        importlib.reload(module)
    type_checked_module_namespace = dict(vars(module))

    # resetting to "default" mode
    importlib.reload(module)

    return {name
            for name, content in type_checked_module_namespace.items()
            if name not in default_module_names
            and inspect.ismodule(content)}

测试

对于包含内容的test.py模块

from typing import TYPE_CHECKING

if TYPE_CHECKING:
    import abc
    from django.utils import timezone

if 'aaa':
    print('bbb')

print('hello world')

给我们

>>> detect_type_checking_mode_modules_names('test')
{'abc', 'timezone'}

检测条件导入的模块

1 个答案:

测试