Lib/knee.py

"""An Python re-implementation of hierarchical module import.

This code is intended to be read, not executed.  However, it does work
-- all you need to do to enable it is "import knee".

(The name is a pun on the klunkier predecessor of this module, "ni".)

"""

import sys, imp, __builtin__, string


# Replacement for __import__()
def import_hook(name, globals=None, locals=None, fromlist=None):
    parent = determine_parent(globals)
    q, tail = find_head_package(parent, name)
    m = load_tail(q, tail)
    if not fromlist:
	return q
    if hasattr(m, "__path__"):
	ensure_fromlist(m, fromlist)
    return m

def determine_parent(globals):
    if not globals or  not globals.has_key("__name__"):
	return None
    pname = globals['__name__']
    if globals.has_key("__path__"):
	parent = sys.modules[pname]
	assert globals is parent.__dict__
	return parent
    if '.' in pname:
	i = string.rfind(pname, '.')
	pname = pname[:i]
	parent = sys.modules[pname]
	assert parent.__name__ == pname
	return parent
    return None

def find_head_package(parent, name):
    if '.' in name:
	i = string.find(name, '.')
	head = name[:i]
	tail = name[i+1:]
    else:
	head = name
	tail = ""
    if parent:
	qname = "%s.%s" % (parent.__name__, head)
    else:
	qname = head
    q = import_module(head, qname, parent)
    if q: return q, tail
    if parent:
	qname = head
	parent = None
	q = import_module(head, qname, parent)
	if q: return q, tail
    raise ImportError, "No module named " + qname

def load_tail(q, tail):
    m = q
    while tail:
	i = string.find(tail, '.')
	if i < 0: i = len(tail)
	head, tail = tail[:i], tail[i+1:]
	mname = "%s.%s" % (m.__name__, head)
	m = import_module(head, mname, m)
	if not m:
	    raise ImportError, "No module named " + mname
    return m

def ensure_fromlist(m, fromlist, recursive=0):
    for sub in fromlist:
	if sub == "*":
	    if not recursive:
		try:
		    all = m.__all__
		except AttributeError:
		    pass
		else:
		    ensure_fromlist(m, all, 1)
	    continue
	if sub != "*" and not hasattr(m, sub):
	    subname = "%s.%s" % (m.__name__, sub)
	    submod = import_module(sub, subname, m)
	    if not submod:
		raise ImportError, "No module named " + subname

def import_module(partname, fqname, parent):
    try:
	return sys.modules[fqname]
    except KeyError:
	pass
    try:
	fp, pathname, stuff = imp.find_module(partname,
					      parent and parent.__path__)
    except ImportError:
	return None
    try:
	m = imp.load_module(fqname, fp, pathname, stuff)
    finally:
	if fp: fp.close()
    if parent:
	setattr(parent, partname, m)
    return m


# Replacement for reload()
def reload_hook(module):
    name = module.__name__
    if '.' not in name:
	return import_module(name, name, None)
    i = string.rfind(name, '.')
    pname = name[:i]
    parent = sys.modules[pname]
    return import_module(name[i+1:], name, parent)


# Save the original hooks
original_import = __builtin__.__import__
original_reload = __builtin__.reload

# Now install our hooks
__builtin__.__import__ = import_hook
__builtin__.reload = reload_hook