如何在运行中替换导入的函数

在公司的一个内部分享会上,同事分享了一个很好的可以在函数执行时,替换导入的函数的方法,让我这等土八路实在是羞愧。
废话不多说,待我娓娓道来。

事件起因

我们写了一个短信通道,为了测试该通道的并发能力,需要进行压力测试,但是压力测试的时候肯定不能真的发短信。之前的办法是在代码里进行一些判断,比如:
if config['DEBUG'] = True:
    dummy_sender(msg)

但是这样很不优雅,而且为了测试脚本,这段代码也要带到生产环境里去,也会影响效率。

Python的特性

Python中一切皆对象,函数也不例外。当我们遇到的是下面的这个脚本时:
from a import sender

def send_sms(msg):
    sender(msg)
Python解释器在调用sender时,会按照sender这个字符串去该脚本中的globals中寻找函数,并且调用。这个脚本中的globals为:
'send_sms': <function send_sms at 0x10a1e7050>,
'__builtins__': <module '__builtin__' (built-in)>,
'__file__': 'practise.py',
'sender': <function sender at 0x10a1e7050>,
那么我们能不能在执行当前脚本的时候,替换掉sender呢。当然可以。send_sms的globals会带有该脚本的所有globals的属性,我们可以从这里下手
send_sms.__globals__['sender'] = dummy_sender
如果打印一下sys.modules, 我们发现,我们函数名称也在里面。于是也可以替换
sys.modules['sender'] = dummy_sender

注意,从a脚本导入后,a脚本导入的函数也会在sys.modules里。
我对于其中的原理虽然了解一些,但还是难以说清楚,按照土八路的思想,能用就万岁。原理以后再补。