如何即时修改导入的源代码?

问题描述

假设我有一个像这样的模块文件:

Suppose I have a module file like this:

# my_module.py
print("hello")

然后我有一个简单的脚本:

Then I have a simple script:

# my_script.py
import my_module

这将打印"hello".

比方说，我想覆盖" print()函数，以便它改为返回"world".我该如何以编程方式执行此操作(无需手动修改my_module.py)?

Let's say I want to "override" the print() function so it returns "world" instead. How could I do this programmatically (without manually modifying my_module.py)?

我想我需要在导入之前或同时以某种方式修改my_module的源代码.显然，导入后无法执行此操作，因此无法使用unittest.mock解决方案.

What I thought is that I need somehow to modify the source code of my_module before or while importing it. Obvisouly, I cannot do this after importing it so solution using unittest.mock are impossible.

我还认为我可以读取文件my_module.py，进行修改，然后加载它.但这很丑陋，因为如果模块位于其他位置，它将无法正常工作.

I also thought I could read the file my_module.py, perform modification, then load it. But this is ugly, as it will not work if the module is located somewhere else.

我认为，好的解决方案是利用 importlib .

The good solution, I think, is to make use of importlib.

我阅读了文档，发现了一个非常交叉的方法:get_source(fullname).我以为我可以覆盖它:

I read the doc and found a very intersecting method: get_source(fullname). I thought I could just override it:

def get_source(fullname):
    source = super().get_source(fullname)
    source = source.replace("hello", "world")
    return source

不幸的是，所有这些抽象类让我有些迷茫，我不知道如何正确执行此操作.

Unfortunately, I am a bit lost with all these abstract classes and I do not know how to perform this properly.

我徒劳地尝试:

spec = importlib.util.find_spec("my_module")
spec.loader.get_source = mocked_get_source
module = importlib.util.module_from_spec(spec)

欢迎任何帮助.

推荐答案

以下是我根据这篇精彩的演讲.允许在导入指定模块之前对源进行任何任意修改.只要幻灯片没有忽略任何重要内容，它就应该是合理正确的.这仅适用于Python 3.5 +.

Here's a solution I hacked together based on the content of this great talk. It allows any arbitrary modifications to be made to the source before importing the specified module. It should be reasonably correct as long as the slides did not omit anything important. This will only work on Python 3.5+.

import importlib
import sys

def modify_and_import(module_name, package, modification_func):
    spec = importlib.util.find_spec(module_name, package)
    source = spec.loader.get_source(module_name)
    new_source = modification_func(source)
    module = importlib.util.module_from_spec(spec)
    codeobj = compile(new_source, module.__spec__.origin, 'exec')
    exec(codeobj, module.__dict__)
    sys.modules[module_name] = module
    return module

因此，您可以使用它

my_module = modify_and_import("my_module", None, lambda src: src.replace("hello", "world"))

这篇关于如何即时修改导入的源代码?的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持！