На этом шаге мы рассмотрим особенности решения этой задачи.

Задача

    Вы хотите изменить функции в существующем модуле или применить к ним декораторы. Однако вы хотите сделать это только в том случае, когда модуль импортируется и где-то используется.

Решение

    Основная проблема тут в том, что вы хотите выполнять действия в ответ на загрузку модуля. Возможно, вы хотите запускать некую функцию обратного вызова, которая будет уведомлять вас о загрузке модуля.

    Эта задача может быть решена с помощью механизма хуков импортирования, который мы обсудили на прошлых шагах. Вот возможное решение:

# postimport.py

import importlib
import sys
from collections import defaultdict

_post_import_hooks = defaultdict(list)


class PostImportFinder:
    def __init__(self):
        self._skip = set()

    def find_module(self, fullname, path=None):
        if fullname in self._skip:
            return None
        self._skip.add(fullname)
        return PostImportLoader(self)


class PostImportLoader:
    def __init__(self, finder):
        self._finder = finder

    def load_module(self, fullname):
        importlib.import_module(fullname)
        module = sys.modules[fullname]
        for func in _post_import_hooks[fullname]:
            func(module)
        self._finder._skip.remove(fullname)
        return module


def when_imported(fullname):
    def decorate(func):
        if fullname in sys.modules:
            func(sys.modules[fullname])
        else:
            _post_import_hooks[fullname].append(func)
        return func

    return decorate


sys.meta_path.insert(0, PostImportFinder())

    Чтобы использовать этот код, примените декоратор when_imported(). Например:

from postimport import when_imported

@when_imported('threading')
def warn_threads(mod):
    print('Threads? Are you crazy?')

if __name__ == '__main__':
    import threading

    Вот результат его выполнения:

Threads? Are you crazy?

    Более практический пример - применение декоратора к существующим определениям:

from postimport import when_imported
from functools import wraps


def logged(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        print('Calling', func.__name__, args, kwargs)
        return func(*args, **kwargs)
    return wrapper

# Пример
@when_imported('math')
def add_logging(mod):
    mod.cos = logged(mod.cos)
    mod.sin = logged(mod.sin)

if __name__ == '__main__':
    import math
    print(math.cos(2))
    print(math.sin(2))

    Вот результат его выполнения:

Calling cos (2,) {}
-0.4161468365471424
Calling sin (2,) {}
0.9092974268256817

Обсуждение

    Этот рецепт опирается на хуки импортирования, которые обсуждаются в прошлых шагах, но с небольшим вывертом.

    Во-первых, роль декоратора @when_imported состоит в том, чтобы зарегистрировать функции-обработчики, которые запускаются при импортировании. Декоратор проверяет sys.modules, чтобы посмотреть, не загружен ли уже модуль. Если это так, обработчик вызывается сразу же. Иначе обработчик будет добавлен в список в словаре _post_import_hooks. Назначение _post_import_hooks в том, чтобы просто собрать все объекты-обработчики, которые зарегистрированы для каждого модуля. В принципе, для передаваемого модуля может быть зарегистрировано больше одного обработчика.

    Чтобы запустить подготовленные действия в _post_import_hooks после импортирования модуля, класс PostImportFinder инсталлируется в качестве первого элемента в sys.meta_path. Как вы помните, sys.meta_path содержит список объектов-поисковиков, к которым по порядку происходят обращения, чтобы найти модули. Путем установки PostImportFinder в качестве первого элемента мы "отлавливаем" все случаи импортирования модулей.

    В этом рецепте, однако, роль PostImportFinder заключается не в загрузке модулей, а в запуске нужных действий по завершении импортирования. Чтобы это сделать, импортирование делегируется другим поисковикам в sys.meta_path. Вместо того чтобы пытаться сделать это напрямую, в классе PostImportLoader рекурсивно вызывается функция imp.import_module(). Чтобы избежать застревания в бесконечном цикле, PostImportFinder хранит набор (множество) всех модулей, которые в настоящий момент находятся в процессе загрузки. Если имя модуля присутствует в этом списке, оно просто игнорируется PostImportFinder. Это вызывает передачу запроса на импортирование другому поисковику из sys.meta_path.

    После того как модуль был загружен с помощью imp.import_module(), все обработчики, в настоящий момент зарегистрированные в post_import_hooks, вызываются с новым загруженным модулем в качестве аргумента. Начиная с этой точки, обработчики свободны делать с модулем все, что пожелают.

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

    Предостережение: этот рецепт не работает для модулей, которые явно загружаются с помощью imp.reload(). Так что если вы перезагружаете ранее загруженный модуль, "послеимпортная" функция-обработчик не запускается заново (еще один довод не использовать reload() в продакшн-коде). С другой стороны, если вы удалите модуль из sys.modules и повторите импорт, то увидите, что обработчик запускается.

    Дополнительную информацию о послеимпортных хуках вы найдете в PEP 369.

    На следующем шаге мы рассмотрим установку пакетов "только для себя".

Предыдущий шаг Содержание Следующий шаг