Making a copy of an entire namespace?

Question

I\'d like to make a copy of an entire namespace while replacing some functions with dynamically constructed versions.

In other words, starting with namespace (impo

Answer 1

To patch a set of functions while importing second instances of a set of functions, you can override the standard Python import hook and apply the patches directly at import time. This will make sure that no other module will ever see the unpatched versions of any of the modules, so even if they import functions from another module directly by name, they will only see the patched functions. Here is a proof-of-concept implementation:

import __builtin__
import collections
import contextlib
import sys


@contextlib.contextmanager
def replace_import_hook(new_import_hook):
    original_import = __builtin__.__import__
    __builtin__.__import__ = new_import_hook
    yield original_import
    __builtin__.__import__ = original_import


def clone_modules(patches, additional_module_names=None):
    """Import new instances of a set of modules with some objects replaced.

    Arguments:
      patches - a dictionary mapping `full.module.name.symbol` to the new object.
      additional_module_names - a list of the additional modules you want new instances of, without
          replacing any objects in them.

    Returns:
      A dictionary mapping module names to the new patched module instances.
    """

    def import_hook(module_name, *args):
        result = original_import(module_name, *args)
        if module_name not in old_modules or module_name in new_modules:
            return result
        # The semantics for the return value of __import__() are a bit weird, so we need some logic
        # to determine the actual imported module object.
        if len(args) >= 3 and args[2]:
            module = result
        else:
            module = reduce(getattr, module_name.split('.')[1:], result)
        for symbol, obj in patches_by_module[module_name].items():
            setattr(module, symbol, obj)
        new_modules[module_name] = module
        return result

    # Group patches by module name
    patches_by_module = collections.defaultdict(dict)
    for dotted_name, obj in patches.items():
        module_name, symbol = dotted_name.rsplit('.', 1)  # Only allows patching top-level objects
        patches_by_module[module_name][symbol] = obj

    try:
        # Remove the old module instances from sys.modules and store them in old_modules
        all_module_names = list(patches_by_module)
        if additional_module_names is not None:
            all_module_names.extend(additional_module_names)
        old_modules = {}
        for name in all_module_names:
            old_modules[name] = sys.modules.pop(name)

        # Re-import modules to create new patched versions
        with replace_import_hook(import_hook) as original_import:
            new_modules = {}
            for module_name in all_module_names:
                import_hook(module_name)
    finally:
        sys.modules.update(old_modules)
    return new_modules

And here some test code for this implementation:

from __future__ import print_function

import math
import random

def patched_log(x):
    print('Computing log({:g})'.format(x))
    return math.log(x)

patches = {'math.log': patched_log}
cloned_modules = clone_modules(patches, ['random'])
new_math = cloned_modules['math']
new_random = cloned_modules['random']
print('Original log:         ', math.log(2.0))
print('Patched log:          ', new_math.log(2.0))
print('Original expovariate: ', random.expovariate(2.0))
print('Patched expovariate:  ', new_random.expovariate(2.0))

The test code has this output:

Computing log(4)
Computing log(4.5)
Original log:          0.69314718056
Computing log(2)
Patched log:           0.69314718056
Original expovariate:  0.00638038735379
Computing log(0.887611)
Patched expovariate:   0.0596108277801

The first two lines of output result from these two lines in random, which are executed at import time. This demonstrates that random sees the patched function right away. The rest of the output demonstrates that the original math and random still use the unpatched version of log, while the cloned modules both use the patched version.

A cleaner way of overriding the import hook might be to use a meta import hook as defined in PEP 302, but providing a full implementation of that approach is beyond the scope of StackOverflow.