时间戳

首页 > 解决方案 > 仅当 Python 模块被导入时才重写它

python - 仅当 Python 模块被导入时才重写它

问题描述

我正在编写一个使用抽象语法树来重写模块部分的库。重写后，我将其放入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.

标签： pythonpython-importabstract-syntax-tree

解决方案

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

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

推荐阅读