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

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

根据这篇精彩演讲的内容，这是我共同构思的解决方案。允许在导入指定模块之前对源进行任何任意修改。只要幻灯片没有忽略任何重要内容，它应该是合理正确的。这仅适用于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"))