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