Lazily import names in a module

#!/usr/bin/env python
# -*-coding: utf8-*-
# Title: lazynames.py
# Author: Gribouillis for the python forum at www.daniweb.com
# Created: 2012-09-25 12:20:15.556784 (isoformat date)
# License: Public Domain
# Use this code freely.

""" module: lazynames -

    Usage:

        This module allows a python module to import symbols lazily
        by declaring a "name getter function" like module A below:
            
            # module A
            import lazynames

            @lazynames.set_getter(__name__)
            def func(module, attribute):
                return attribute.upper()
        
        The client module B below imports symbols from A which
        are not defined in A 
        
            # module B
            from A import foo, bar, baz
            print(foo, bar, baz) # prints FOO BAR BAZ
            
        When foo, bar, baz are imported, module A uses its
        "name getter function" to create values for these names.
        Optionaly, the values can be stored in module A, so that
        they would be created only once:
            
            @lazynames.set_getter(__name__)
            def func(module, attribute):
                value = attribute.upper()
                setattr(module, attribute, value)
                return value
        
        The following operations are possible:
            
            lazynames.set_getter("my.mod")(func)
                Sets a name getter function for the module named "my.mod".
                Module my.mod must already exist in sys.modules for this
                to work.
                This can be used more than once for the same module.
                lazynames.set_getter("my.mod") can also be used as a
                function decorator as shown above.
                The "name getter function" can be called __getattr__ as
                it behaves like a __getattr__ function: an expression
                like my.mod.foo calls the name getter function with
                argument "foo" if the name "foo" does not already exist
                in module my.mod. The name getter function should raise
                AttributeError to indicate that it can not create a value
                for a given name.
                By default, set_getter won't set a name getter for the
                __main__ module, unless a special flag is passed:
                    lazynames.set_getter("__main__", allow_main = True)(func)
                This means that if the example module A above is used
                as a script instead of being imported, the name getter
                won't be set.
            
            lazynames.get_getter("my.mod")
                returns the current name getter function of module my.mod
                if it has been defined, otherwise return None

"""

from collections import namedtuple
from functools import partial
import sys

__all__ = ["get_getter", "set_getter", "LazyModule"]

Record = namedtuple("Record", "module getter")
REGISTERED = dict()

class LazyModule(object):
    def __init__(self, dict):
        self.__dict__ = dict
        
    def __getattr__(self, attr):
        record = REGISTERED[self.__name__]
        try:
            return getattr(record.module, attr)
        except AttributeError:
            if attr.startswith("__"):
                raise
            else:
                return record.getter(self, attr)

def get_getter(name):
    return REGISTERED[name].getter if name in REGISTERED else None

def set_getter(name, allow_main = False):
    return partial(_register, name = name, allow_main = allow_main)

def _register(getter_func, name = "__main__", allow_main = False):
    if getter_func is None:
        _unregister(name)
    elif allow_main or name != "__main__":
        mod = (REGISTERED[name].module if name in REGISTERED
                else sys.modules[name])
        record = Record(module = mod, getter = getter_func)
        REGISTERED[name] = record
        sys.modules[name] = LazyModule(mod.__dict__)
    return getter_func

def _unregister(name):
    try:
        record = REGISTERED.pop(name)
    except KeyError:
        pass
    else:
        sys.modules[name] = record.module