拦截函数调用的Pythonic方式？

Question

在测试一些需要查询环境的东西时，比如说 os.getenv 或 sys.version 等，直接让查询结果“说谎”往往比真的伪造环境要方便得多。下面是一个上下文管理器，它可以一次性处理一个 os.getenv 的调用：

from __future__ import with_statement
from contextlib import contextmanager
import os

@contextmanager
def fake_env(**fakes):
    '''fakes is a dict mapping variables to their values. In the
    fake_env context, os.getenv calls try to return out of the fakes
    dict whenever possible before querying the actual environment.
    '''

    global os
    original = os.getenv

    def dummy(var):
        try: return fakes[var]
        except KeyError: return original(var)

    os.getenv = dummy
    yield
    os.getenv = original

if __name__ == '__main__':

    print os.getenv('HOME') 
    with fake_env(HOME='here'):
        print os.getenv('HOME') 
    print os.getenv('HOME') 

不过，这个方法只适用于 os.getenv，而且如果我想处理多个参数的函数，语法就会变得有点复杂。我想通过 ast 和 code/exec/eval 来扩展这个功能，让它可以接收要覆盖的函数作为参数，但这样做不够简洁。而且，这样一来，我就可能会走上 Greenspun 的第十条法则。有没有更好的方法呢？

Answer 1

为什么不自己写一个（假的）sys、os等模块呢？

import fakeSys as sys

Answer 2

你可以很简单地把 os.getenv 直接作为第一个参数传进去，然后在上下文管理器里分析它，这样比用 ast、code 等等要简单得多：

>>> os.getenv.__name__
'getenv'
>>> os.getenv.__module__
'os'

之后，为了更通用的使用，你可以返回一个结果对象，或者把参数（可能是元组）映射到结果上。这个 faker 上下文管理器也可以选择性地接受一个可调用的函数来进行伪造。

举个例子，简单到极致：

import sys

def faker(original, fakefun):

    original = os.getenv
    themod = sys.modules[original.__module__]
    thename = original.__name__

    def dummy(*a, **k):
        try: return fakefun(*a, **k)
        except BaseException: return original(*a, **k)

    setattr(themod, thename, dummy)
    yield
    setattr(themod, thename, original)

你的具体例子可以变成：

with faker(os.getenv, dict(HOME='here').__getitem__):
   ...

当然，如果你想传播某些异常，而不是直接返回原始函数的结果，或者在某些常见情况下提供伪造函数比较麻烦，那么稍微复杂一点也是合理的。但没有理由让这样一个通用的伪造器比你具体的那个复杂得多。