Fixed issue where patcher wasn't correctly restoring modules after an importerror.

This commit is contained in:
Ryan Williams
2010-01-16 00:34:19 -08:00
parent bf36075800
commit 4a15b5b81c
2 changed files with 23 additions and 21 deletions

View File

@@ -13,28 +13,30 @@ def inject(module_name, new_globals, *additional_modules):
## Remove the old module from sys.modules and reimport it while the specified modules are in place ## Remove the old module from sys.modules and reimport it while the specified modules are in place
old_module = sys.modules.pop(module_name, None) old_module = sys.modules.pop(module_name, None)
module = __import__(module_name, {}, {}, module_name.split('.')[:-1]) try:
module = __import__(module_name, {}, {}, module_name.split('.')[:-1])
if new_globals is not None: if new_globals is not None:
## Update the given globals dictionary with everything from this new module ## Update the given globals dictionary with everything from this new module
for name in dir(module): for name in dir(module):
if name not in __exclude: if name not in __exclude:
new_globals[name] = getattr(module, name) new_globals[name] = getattr(module, name)
## Keep a reference to the new module to prevent it from dying ## Keep a reference to the new module to prevent it from dying
sys.modules['__patched_module_' + module_name] = module sys.modules['__patched_module_' + module_name] = module
## Put the original module back finally:
if old_module is not None: ## Put the original module back
sys.modules[module_name] = old_module if old_module is not None:
else: sys.modules[module_name] = old_module
del sys.modules[module_name] elif module_name in sys.modules:
del sys.modules[module_name]
## Put all the saved modules back ## Put all the saved modules back
for name, mod in additional_modules: for name, mod in additional_modules:
if saved[name] is not None: if saved[name] is not None:
sys.modules[name] = saved[name] sys.modules[name] = saved[name]
else: else:
del sys.modules[name] del sys.modules[name]
return module return module