仅在导入时重写 Python 模块

Question

我正在编写一个库，它使用抽象语法树来重写模块。重写后，我将其放在 sys.modules 中，以便其他模块可以叫它。但是，时机很重要，我不能只是运行重写的模块在开始时。我希望它在被另一个模块导入时为运行，而不是之前。

我已经通过编写一个 importer 解决了这个问题，但是它使用了 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.

Answer 1

看来我误解了进口协议。您可以覆盖执行模块的部分，并保留创建新模块的部分不变。这是我的示例，重写为使用更新的导入器协议 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()

输出与旧版本完全相同，但弃用警告消失了。

仅在导入时重写 Python 模块

Rewriting a Python module only if it gets imported

python

abstract-syntax-tree

python-import