Pythonic way to intercept calls to a function?

Question

For testing things that query the environment (e.g., os.getenv, sys.version, etc.), it's often more convenient to make the queries lie than to actually fake up the environment. Here's a context manager that does this for one os.getenv call at a time:

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') 

But this only works for os.getenv and the syntax gets a bit clunky if I allow for functions with multiple arguments. I guess between ast and code/exec/eval I could extend it to take the function to override as a parameter, but not cleanly. Also, I would then be on my way to Greenspun's Tenth. Is there a better way?

Answer 1

Why not write your own (fake) sys, os, &c. modules?

import fakeSys as sys

Answer 2

You could easily just pass os.getenv itself as the first argument, then analyze it in the context manager much more simply than ast, code, etc etc:

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

After that, for reasonably general purpose use, you could have the result object to be returned, or a mapping from arguments (probably tuples thereof) to results. The faker context manager could also optionally accept a callable to be used for faking.

For example, with maximum simplicity:

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)

Your specific example could become:

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

Of course, a little more complexity may be warranted if e.g. you want to propagate certain exceptions rather than punting to the original function, or shortcut some common cases where providing a fakefun callable is clunky, and so on. But there's no reason such a general faker need be much more complex than your specific one.