smoketests/sin_testing/runpy_backport.py

   1 """runpy.py - locating and running Python code using the module namespace
   2
   3 Provides support for locating and running Python scripts using the Python
   4 module namespace instead of the native filesystem.
   5
   6 This allows Python code to play nicely with non-filesystem based PEP 302
   7 importers when locating support scripts as well as when importing modules.
   8 """
   9
  10 # Written by Nick Coghlan <ncoghlan at gmail.com>
  11 #    to implement PEP 338 (Executing Modules as Scripts)
  12
  13 import sys
  14 import imp
  15 from pkgutil import read_code
  16 try:
  17     from imp import get_loader
  18 except ImportError:
  19     from pkgutil import get_loader
  20
  21 __all__ = [
  22     "run_module", "run_path",
  23 ]
  24
  25 class _TempModule(object):
  26     """Temporarily replace a module in sys.modules with an empty namespace"""
  27     def __init__(self, mod_name):
  28         self.mod_name = mod_name
  29         self.module = imp.new_module(mod_name)
  30         self._saved_module = []
  31
  32     def __enter__(self):
  33         mod_name = self.mod_name
  34         try:
  35             self._saved_module.append(sys.modules[mod_name])
  36         except KeyError:
  37             pass
  38         sys.modules[mod_name] = self.module
  39         return self
  40
  41     def __exit__(self, *args):
  42         if self._saved_module:
  43             sys.modules[self.mod_name] = self._saved_module[0]
  44         else:
  45             del sys.modules[self.mod_name]
  46         self._saved_module = []
  47
  48 class _ModifiedArgv0(object):
  49     def __init__(self, value):
  50         self.value = value
  51         self._saved_value = self._sentinel = object()
  52
  53     def __enter__(self):
  54         if self._saved_value is not self._sentinel:
  55             raise RuntimeError("Already preserving saved value")
  56         self._saved_value = sys.argv[0]
  57         sys.argv[0] = self.value
  58
  59     def __exit__(self, *args):
  60         self.value = self._sentinel
  61         sys.argv[0] = self._saved_value
  62
  63 def _run_code(code, run_globals, init_globals=None,
  64               mod_name=None, mod_fname=None,
  65               mod_loader=None, pkg_name=None):
  66     """Helper to run code in nominated namespace"""
  67     if init_globals is not None:
  68         run_globals.update(init_globals)
  69     run_globals.update(__name__ = mod_name,
  70                        __file__ = mod_fname,
  71                        __loader__ = mod_loader,
  72                        __package__ = pkg_name)
  73     exec code in run_globals
  74     return run_globals
  75
  76 def _run_module_code(code, init_globals=None,
  77                     mod_name=None, mod_fname=None,
  78                     mod_loader=None, pkg_name=None):
  79     """Helper to run code in new namespace with sys modified"""
  80     with _TempModule(mod_name) as temp_module:
  81         _ModifiedArgv0(mod_fname),
  82         mod_globals = temp_module.module.__dict__
  83         _run_code(code, mod_globals, init_globals,
  84                   mod_name, mod_fname, mod_loader, pkg_name)
  85     # Copy the globals of the temporary module, as they
  86     # may be cleared when the temporary module goes away
  87     return mod_globals.copy()
  88
  89
  90 # This helper is needed due to a missing component in the PEP 302
  91 # loader protocol (specifically, "get_filename" is non-standard)
  92 # Since we can't introduce new features in maintenance releases,
  93 # support was added to zipimporter under the name '_get_filename'
  94 def _get_filename(loader, mod_name):
  95     for attr in ("get_filename", "_get_filename"):
  96         meth = getattr(loader, attr, None)
  97         if meth is not None:
  98             return meth(mod_name)
  99     return None
 100
 101 # Helper to get the loader, code and filename for a module
 102 def _get_module_details(mod_name):
 103     loader = get_loader(mod_name)
 104     if loader is None:
 105         raise ImportError("No module named %s" % mod_name)
 106     if loader.is_package(mod_name):
 107         if mod_name == "__main__" or mod_name.endswith(".__main__"):
 108             raise ImportError("Cannot use package as __main__ module")
 109         try:
 110             pkg_main_name = mod_name + ".__main__"
 111             return _get_module_details(pkg_main_name)
 112         except ImportError, e:
 113             raise ImportError(("%s; %r is a package and cannot " +
 114                                "be directly executed") %(e, mod_name))
 115     code = loader.get_code(mod_name)
 116     if code is None:
 117         raise ImportError("No code object available for %s" % mod_name)
 118     filename = _get_filename(loader, mod_name)
 119     return mod_name, loader, code, filename
 120
 121
 122 def _get_main_module_details():
 123     # Helper that gives a nicer error message when attempting to
 124     # execute a zipfile or directory by invoking __main__.py
 125     main_name = "__main__"
 126     try:
 127         return _get_module_details(main_name)
 128     except ImportError as exc:
 129         if main_name in str(exc):
 130             raise ImportError("can't find %r module in %r" %
 131                               (main_name, sys.path[0]))
 132         raise
 133
 134 # This function is the actual implementation of the -m switch and direct
 135 # execution of zipfiles and directories and is deliberately kept private.
 136 # This avoids a repeat of the situation where run_module() no longer met the
 137 # needs of mainmodule.c, but couldn't be changed because it was public
 138 def _run_module_as_main(mod_name, alter_argv=True):
 139     """Runs the designated module in the __main__ namespace
 140
 141        Note that the executed module will have full access to the
 142        __main__ namespace. If this is not desirable, the run_module()
 143        function should be used to run the module code in a fresh namespace.
 144
 145        At the very least, these variables in __main__ will be overwritten:
 146            __name__
 147            __file__
 148            __loader__
 149            __package__
 150     """
 151     try:
 152         if alter_argv or mod_name != "__main__": # i.e. -m switch
 153             mod_name, loader, code, fname = _get_module_details(mod_name)
 154         else:          # i.e. directory or zipfile execution
 155             mod_name, loader, code, fname = _get_main_module_details()
 156     except ImportError as exc:
 157         msg = "%s: %s" % (sys.executable, str(exc))
 158         sys.exit(msg)
 159     pkg_name = mod_name.rpartition('.')[0]
 160     main_globals = sys.modules["__main__"].__dict__
 161     if alter_argv:
 162         sys.argv[0] = fname
 163     return _run_code(code, main_globals, None,
 164                      "__main__", fname, loader, pkg_name)
 165
 166 def run_module(mod_name, init_globals=None,
 167                run_name=None, alter_sys=False):
 168     """Execute a module's code without importing it
 169
 170        Returns the resulting top level namespace dictionary
 171     """
 172     mod_name, loader, code, fname = _get_module_details(mod_name)
 173     if run_name is None:
 174         run_name = mod_name
 175     pkg_name = mod_name.rpartition('.')[0]
 176     if alter_sys:
 177         return _run_module_code(code, init_globals, run_name,
 178                                 fname, loader, pkg_name)
 179     else:
 180         # Leave the sys module alone
 181         return _run_code(code, {}, init_globals, run_name,
 182                          fname, loader, pkg_name)
 183
 184
 185 # XXX (ncoghlan): Perhaps expose the C API function
 186 # as imp.get_importer instead of reimplementing it in Python?
 187 def _get_importer(path_name):
 188     """Python version of PyImport_GetImporter C API function"""
 189     cache = sys.path_importer_cache
 190     try:
 191         importer = cache[path_name]
 192     except KeyError:
 193         # Not yet cached. Flag as using the
 194         # standard machinery until we finish
 195         # checking the hooks
 196         cache[path_name] = None
 197         for hook in sys.path_hooks:
 198             try:
 199                 importer = hook(path_name)
 200                 break
 201             except ImportError:
 202                 pass
 203         else:
 204             # The following check looks a bit odd. The trick is that
 205             # NullImporter throws ImportError if the supplied path is a
 206             # *valid* directory entry (and hence able to be handled
 207             # by the standard import machinery)
 208             try:
 209                 importer = imp.NullImporter(path_name)
 210             except ImportError:
 211                 return None
 212         cache[path_name] = importer
 213     return importer
 214
 215 def _get_code_from_file(fname):
 216     # Check for a compiled file first
 217     with open(fname, "rb") as f:
 218         code = read_code(f)
 219     if code is None:
 220         # That didn't work, so try it as normal source code
 221         with open(fname, "rU") as f:
 222             code = compile(f.read(), fname, 'exec')
 223     return code
 224
 225 def run_path(path_name, init_globals=None, run_name=None):
 226     """Execute code located at the specified filesystem location
 227
 228        Returns the resulting top level namespace dictionary
 229
 230        The file path may refer directly to a Python script (i.e.
 231        one that could be directly executed with execfile) or else
 232        it may refer to a zipfile or directory containing a top
 233        level __main__.py script.
 234     """
 235     if run_name is None:
 236         run_name = "<run_path>"
 237     importer = _get_importer(path_name)
 238     if isinstance(importer, imp.NullImporter):
 239         # Not a valid sys.path entry, so run the code directly
 240         # execfile() doesn't help as we want to allow compiled files
 241         code = _get_code_from_file(path_name)
 242         return _run_module_code(code, init_globals, run_name, path_name)
 243     else:
 244         # Importer is defined for path, so add it to
 245         # the start of sys.path
 246         sys.path.insert(0, path_name)
 247         try:
 248             # Here's where things are a little different from the run_module
 249             # case. There, we only had to replace the module in sys while the
 250             # code was running and doing so was somewhat optional. Here, we
 251             # have no choice and we have to remove it even while we read the
 252             # code. If we don't do this, a __loader__ attribute in the
 253             # existing __main__ module may prevent location of the new module.
 254             main_name = "__main__"
 255             saved_main = sys.modules[main_name]
 256             del sys.modules[main_name]
 257             try:
 258                 mod_name, loader, code, fname = _get_main_module_details()
 259             finally:
 260                 sys.modules[main_name] = saved_main
 261             pkg_name = ""
 262             with _TempModule(run_name) as temp_module:
 263                 _ModifiedArgv0(path_name),
 264                 mod_globals = temp_module.module.__dict__
 265                 return _run_code(code, mod_globals, init_globals,
 266                                     run_name, fname, loader, pkg_name).copy()
 267         finally:
 268             try:
 269                 sys.path.remove(path_name)
 270             except ValueError:
 271                 pass
 272
 273
 274 if __name__ == "__main__":
 275     # Run the module specified as the next command line argument
 276     if len(sys.argv) < 2:
 277         print >> sys.stderr, "No module specified for execution"
 278     else:
 279         del sys.argv[0] # Make the requested module sys.argv[0]
 280         _run_module_as_main(sys.argv[0])