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):
101    loader = get_loader(mod_name)
102    if loader is None:
103        raise ImportError("No module named %s" % mod_name)
104    if loader.is_package(mod_name):
105        if mod_name == "__main__" or mod_name.endswith(".__main__"):
106            raise ImportError("Cannot use package as __main__ module")
107        try:
108            pkg_main_name = mod_name + ".__main__"
109            return _get_module_details(pkg_main_name)
110        except ImportError, e:
111            raise ImportError(("%s; %r is a package and cannot " +
112                               "be directly executed") %(e, mod_name))
113    code = loader.get_code(mod_name)
114    if code is None:
115        raise ImportError("No code object available for %s" % mod_name)
116    filename = _get_filename(loader, mod_name)
117    return mod_name, loader, code, filename
118
119
120def _get_main_module_details():
121    # Helper that gives a nicer error message when attempting to
122    # execute a zipfile or directory by invoking __main__.py
123    main_name = "__main__"
124    try:
125        return _get_module_details(main_name)
126    except ImportError as exc:
127        if main_name in str(exc):
128            raise ImportError("can't find %r module in %r" %
129                              (main_name, sys.path[0]))
130        raise
131
132# This function is the actual implementation of the -m switch and direct
133# execution of zipfiles and directories and is deliberately kept private.
134# This avoids a repeat of the situation where run_module() no longer met the
135# needs of mainmodule.c, but couldn't be changed because it was public
136def _run_module_as_main(mod_name, alter_argv=True):
137    """Runs the designated module in the __main__ namespace
138
139       Note that the executed module will have full access to the
140       __main__ namespace. If this is not desirable, the run_module()
141       function should be used to run the module code in a fresh namespace.
142
143       At the very least, these variables in __main__ will be overwritten:
144           __name__
145           __file__
146           __loader__
147           __package__
148    """
149    try:
150        if alter_argv or mod_name != "__main__": # i.e. -m switch
151            mod_name, loader, code, fname = _get_module_details(mod_name)
152        else:          # i.e. directory or zipfile execution
153            mod_name, loader, code, fname = _get_main_module_details()
154    except ImportError as exc:
155        msg = "%s: %s" % (sys.executable, str(exc))
156        sys.exit(msg)
157    pkg_name = mod_name.rpartition('.')[0]
158    main_globals = sys.modules["__main__"].__dict__
159    if alter_argv:
160        sys.argv[0] = fname
161    return _run_code(code, main_globals, None,
162                     "__main__", fname, loader, pkg_name)
163
164def run_module(mod_name, init_globals=None,
165               run_name=None, alter_sys=False):
166    """Execute a module's code without importing it
167
168       Returns the resulting top level namespace dictionary
169    """
170    mod_name, loader, code, fname = _get_module_details(mod_name)
171    if run_name is None:
172        run_name = mod_name
173    pkg_name = mod_name.rpartition('.')[0]
174    if alter_sys:
175        return _run_module_code(code, init_globals, run_name,
176                                fname, loader, pkg_name)
177    else:
178        # Leave the sys module alone
179        return _run_code(code, {}, init_globals, run_name,
180                         fname, loader, pkg_name)
181
182
183# XXX (ncoghlan): Perhaps expose the C API function
184# as imp.get_importer instead of reimplementing it in Python?
185def _get_importer(path_name):
186    """Python version of PyImport_GetImporter C API function"""
187    cache = sys.path_importer_cache
188    try:
189        importer = cache[path_name]
190    except KeyError:
191        # Not yet cached. Flag as using the
192        # standard machinery until we finish
193        # checking the hooks
194        cache[path_name] = None
195        for hook in sys.path_hooks:
196            try:
197                importer = hook(path_name)
198                break
199            except ImportError:
200                pass
201        else:
202            # The following check looks a bit odd. The trick is that
203            # NullImporter raises ImportError if the supplied path is a
204            # *valid* directory entry (and hence able to be handled
205            # by the standard import machinery)
206            try:
207                importer = imp.NullImporter(path_name)
208            except ImportError:
209                return None
210        cache[path_name] = importer
211    return importer
212
213def _get_code_from_file(fname):
214    # Check for a compiled file first
215    with open(fname, "rb") as f:
216        code = read_code(f)
217    if code is None:
218        # That didn't work, so try it as normal source code
219        with open(fname, "rU") as f:
220            code = compile(f.read(), fname, 'exec')
221    return code
222
223def run_path(path_name, init_globals=None, run_name=None):
224    """Execute code located at the specified filesystem location
225
226       Returns the resulting top level namespace dictionary
227
228       The file path may refer directly to a Python script (i.e.
229       one that could be directly executed with execfile) or else
230       it may refer to a zipfile or directory containing a top
231       level __main__.py script.
232    """
233    if run_name is None:
234        run_name = "<run_path>"
235    importer = _get_importer(path_name)
236    if isinstance(importer, imp.NullImporter):
237        # Not a valid sys.path entry, so run the code directly
238        # execfile() doesn't help as we want to allow compiled files
239        code = _get_code_from_file(path_name)
240        return _run_module_code(code, init_globals, run_name, path_name)
241    else:
242        # Importer is defined for path, so add it to
243        # the start of sys.path
244        sys.path.insert(0, path_name)
245        try:
246            # Here's where things are a little different from the run_module
247            # case. There, we only had to replace the module in sys while the
248            # code was running and doing so was somewhat optional. Here, we
249            # have no choice and we have to remove it even while we read the
250            # code. If we don't do this, a __loader__ attribute in the
251            # existing __main__ module may prevent location of the new module.
252            main_name = "__main__"
253            saved_main = sys.modules[main_name]
254            del sys.modules[main_name]
255            try:
256                mod_name, loader, code, fname = _get_main_module_details()
257            finally:
258                sys.modules[main_name] = saved_main
259            pkg_name = ""
260            with _TempModule(run_name) as temp_module, \
261                 _ModifiedArgv0(path_name):
262                mod_globals = temp_module.module.__dict__
263                return _run_code(code, mod_globals, init_globals,
264                                    run_name, fname, loader, pkg_name).copy()
265        finally:
266            try:
267                sys.path.remove(path_name)
268            except ValueError:
269                pass
270
271
272if __name__ == "__main__":
273    # Run the module specified as the next command line argument
274    if len(sys.argv) < 2:
275        print >> sys.stderr, "No module specified for execution"
276    else:
277        del sys.argv[0] # Make the requested module sys.argv[0]
278        _run_module_as_main(sys.argv[0])
279