Обнаружение модулей, которые импортируются при условии - PullRequest
Новая
       6

Обнаружение модулей, которые импортируются при условии

0 голосов
/ 19 февраля 2019

Мне нужно извлечь имена модулей, которые импортируются в файл mypy, например: 

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

На основе приведенного выше примера функция "name-extractor" должна вернуть django.utils.timezone и abc, но он не должен возвращать typing.TYPE_CHECKING.

Единственный способ, который я мог бы придумать, - это использование библиотек 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)

Есть ли правильный способ сделать это?

1 Ответ

0 голосов
/ 19 февраля 2019

Ваш первоначальный подход с ast - хорошее начало, но вы можете в конечном итоге пропустить некоторые крайние случаи (например, те, которые я уже упоминал в комментариях), это большая работа, которую можно избежать.

Если вы не возражаете против импорта анализируемого вами модуля (и мы уверены, что присутствуют все подмодули, такие как django), то мы можем сделать это «грязным» способом:

  1. Модуль импорта с TYPE_CHECKING, установленным на False (по умолчанию).
  2. Модуль импорта с TYPE_CHECKING, установленным на True.
  3. Сравните 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)

, после этого наша функция может быть написана с использованием модуля importlib (для динамического импорта и перезагрузки).) и inspect module (для проверки, является ли объект модулем) из stdlib, например 

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

Для 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'}
...