仅在导入时重写Python模块

时间:2018-12-05 06:18:25

标签: python python-import abstract-syntax-tree

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

我已经通过编写importer来解决了这个问题,但是使用了imp模块 为我的重写代码创建一个新的模块对象。该imp模块现在 已弃用,并且替换似乎无法让我创建并执行 新模块。它只是让我找到源文件,并创建一个spec对象 指向那个。

如果我再也不能使用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 :(得分:0)

似乎我误解了导入程序协议。您可以覆盖执行模块的部分,而保留创建新模块的部分。这是我的示例,重写后使用了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()

该版本的输出与旧版本完全相同,但是已弃用警告消失了。