python - 如何即时修改导入的源代码？

Question

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

# my_module.py
print("hello")

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

# my_script.py
import my_module

这将打印"hello"。

假设我想“覆盖”print() 函数，使其返回"world"。我如何以编程方式执行此操作(无需手动修改 my_module.py)？

---

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

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

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

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

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

不幸的是，我对所有这些抽象类有点迷茫，我不知道如何正确地执行它。

我试过了:

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

欢迎任何帮助。

Answer 1

这是一个基于this great talk 内容的解决方案.它允许在导入指定模块之前对源代码进行任意修改。只要幻灯片没有遗漏任何重要内容，它就应该是相当正确的。这仅适用于 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"))