Переписывать модуль Python, только если он импортирован - PullRequest
Новая
       8

Переписывать модуль Python, только если он импортирован

0 голосов
/ 05 декабря 2018

Я пишу библиотеку, которая использует абстрактные синтаксические деревья для переписывания частей модуля.Когда он переписан, я помещаю его в sys.modules, чтобы другие модули могли вызывать его.Тем не менее, время важно, и я не могу просто запустить переписанный модуль в начале.Я хочу, чтобы он запускался, когда он импортируется другим модулем, а не раньше.

Я решил это, написав импортер , но для создания нового использовался модуль impобъект модуля для моего переписанного кода.Этот модуль imp устарел, и его замена не позволяет мне создавать и выполнять новый модуль.Он просто позволяет мне найти исходный файл и создать объект спецификации, указывающий на это.

Если я больше не могу использовать модуль imp, как я могу создать новый модуль с переписанным кодом?

В качестве тривиального примера, у меня есть модуль, который просто распечатывает пару сообщений: 

# my_module.py
print('This is in my_module.py.')

def do_something():
    print('Doing something.')

Мой трассировщик имеет опции для того, чтобы импортировать my_module.py или нет, или нетперепишите его с дополнительным print() сообщением. 

# tracer.py
import builtins
import imp
import sys
from argparse import ArgumentParser
from ast import NodeTransformer, Expr, Call, Name, Load, Str, parse, fix_missing_locations
from pathlib import Path


def main():
    print('Starting.')
    args = parse_args()

    if args.traced:
        sys.meta_path.insert(0, TracedModuleImporter('my_module'))
        print('Set up tracing.')

    if args.imported:
        from my_module import do_something
        do_something()

    print('Done.')


class TracedModuleImporter(object):
    PSEUDO_FILENAME = '<traced>'

    def __init__(self, fullname):
        self.fullname = fullname
        source = Path(fullname + '.py').read_text()
        tree = parse(source, self.PSEUDO_FILENAME)
        new_tree = Tracer().visit(tree)
        fix_missing_locations(new_tree)
        self.code = compile(new_tree, self.PSEUDO_FILENAME, 'exec')

    def find_module(self, fullname, path=None):
        if fullname != self.fullname:
            return None
        return self

    def load_module(self, fullname):
        new_mod = imp.new_module(fullname)
        sys.modules[fullname] = new_mod
        new_mod.__builtins__ = builtins
        new_mod.__file__ = self.PSEUDO_FILENAME
        new_mod.__package__ = None

        exec(self.code, new_mod.__dict__)
        return new_mod


class Tracer(NodeTransformer):
    def visit_Module(self, node):
        new_node = self.generic_visit(node)
        new_node.body.append(Expr(value=Call(func=Name(id='print', ctx=Load()),
                                             args=[Str(s='Traced')],
                                             keywords=[])))
        return new_node


def parse_args():
    parser = ArgumentParser()
    parser.add_argument('--imported', action='store_true')
    parser.add_argument('--traced', action='store_true')
    return parser.parse_args()


main()

Когда я его вызываю, вы можете увидеть сообщения: 

$ python tracer.py
Starting.
Done.
$ python tracer.py --imported
Starting.
This is in my_module.py.
Doing something.
Done.
$ python tracer.py --imported --traced
Starting.
Set up tracing.
This is in my_module.py.
Traced
Doing something.
Done.
$ python tracer.py --traced
Starting.
Set up tracing.
Done.

Все это прекрасно работает с Python 3.6, но Python 3.7жалуется на модуль imp: 

$ python tracer.py
tracer.py:100: DeprecationWarning: the imp module is deprecated in favour of importlib; see the module's documentation for alternative uses
  import imp
Starting.
Done.

1 Ответ

0 голосов
/ 05 декабря 2018

Похоже, я неправильно понял протокол импортера.Вы можете переопределить часть, которая выполняет модуль, и оставить часть, которая создает новый модуль, без изменений.Вот мой пример, переписанный для использования более нового протокола импорта с find_spec() и execute_module() вместо find_module() и load_module(). 

import sys
from argparse import ArgumentParser
from ast import NodeTransformer, Expr, Call, Name, Load, Str, parse, fix_missing_locations
from importlib.abc import MetaPathFinder, Loader
from importlib.machinery import ModuleSpec
from pathlib import Path


def main():
    print('Starting.')
    args = parse_args()

    if args.traced:
        sys.meta_path.insert(0, TracedModuleImporter('my_module'))
        print('Set up tracing.')

    if args.imported:
        from my_module import do_something
        do_something()

    print('Done.')


class TracedModuleImporter(MetaPathFinder, Loader):
    PSEUDO_FILENAME = '<traced>'

    def __init__(self, fullname):
        self.fullname = fullname
        source = Path(fullname + '.py').read_text()
        tree = parse(source, self.PSEUDO_FILENAME)
        new_tree = Tracer().visit(tree)
        fix_missing_locations(new_tree)
        self.code = compile(new_tree, self.PSEUDO_FILENAME, 'exec')

    def find_spec(self, fullname, path, target=None):
        if fullname != self.fullname:
            return None
        return ModuleSpec(fullname, self)

    def exec_module(self, module):
        module.__file__ = self.PSEUDO_FILENAME
        exec(self.code, module.__dict__)


class Tracer(NodeTransformer):
    def visit_Module(self, node):
        new_node = self.generic_visit(node)
        new_node.body.append(Expr(value=Call(func=Name(id='print', ctx=Load()),
                                             args=[Str(s='Traced')],
                                             keywords=[])))
        return new_node


def parse_args():
    parser = ArgumentParser()
    parser.add_argument('--imported', action='store_true')
    parser.add_argument('--traced', action='store_true')
    return parser.parse_args()


main()

Вывод этого точно такой же, как в старой версии, но предупреждение об устаревании пропало.

...