runpy.py 10.2 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11
"""runpy.py - locating and running Python code using the module namespace

Provides support for locating and running Python scripts using the Python
module namespace instead of the native filesystem.

This allows Python code to play nicely with non-filesystem based PEP 302
importers when locating support scripts as well as when importing modules.
"""
# Written by Nick Coghlan <ncoghlan at gmail.com>
#    to implement PEP 338 (Executing Modules as Scripts)

12

13
import os
14
import sys
15
import importlib.machinery # importlib first so we can test #15386 via -m
16
import imp
17
from pkgutil import read_code, get_loader, get_importer
18 19

__all__ = [
Nick Coghlan's avatar
Nick Coghlan committed
20
    "run_module", "run_path",
21 22
]

Nick Coghlan's avatar
Nick Coghlan committed
23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59
class _TempModule(object):
    """Temporarily replace a module in sys.modules with an empty namespace"""
    def __init__(self, mod_name):
        self.mod_name = mod_name
        self.module = imp.new_module(mod_name)
        self._saved_module = []

    def __enter__(self):
        mod_name = self.mod_name
        try:
            self._saved_module.append(sys.modules[mod_name])
        except KeyError:
            pass
        sys.modules[mod_name] = self.module
        return self

    def __exit__(self, *args):
        if self._saved_module:
            sys.modules[self.mod_name] = self._saved_module[0]
        else:
            del sys.modules[self.mod_name]
        self._saved_module = []

class _ModifiedArgv0(object):
    def __init__(self, value):
        self.value = value
        self._saved_value = self._sentinel = object()

    def __enter__(self):
        if self._saved_value is not self._sentinel:
            raise RuntimeError("Already preserving saved value")
        self._saved_value = sys.argv[0]
        sys.argv[0] = self.value

    def __exit__(self, *args):
        self.value = self._sentinel
        sys.argv[0] = self._saved_value
60

61 62
def _run_code(code, run_globals, init_globals=None,
              mod_name=None, mod_fname=None,
63
              mod_loader=None, pkg_name=None):
Benjamin Peterson's avatar
Benjamin Peterson committed
64
    """Helper to run code in nominated namespace"""
65 66 67 68
    if init_globals is not None:
        run_globals.update(init_globals)
    run_globals.update(__name__ = mod_name,
                       __file__ = mod_fname,
Barry Warsaw's avatar
Barry Warsaw committed
69
                       __cached__ = None,
70
                       __doc__ = None,
71 72
                       __loader__ = mod_loader,
                       __package__ = pkg_name)
73
    exec(code, run_globals)
74 75 76
    return run_globals

def _run_module_code(code, init_globals=None,
77
                    mod_name=None, mod_fname=None,
78
                    mod_loader=None, pkg_name=None):
Benjamin Peterson's avatar
Benjamin Peterson committed
79
    """Helper to run code in new namespace with sys modified"""
Nick Coghlan's avatar
Nick Coghlan committed
80 81
    with _TempModule(mod_name) as temp_module, _ModifiedArgv0(mod_fname):
        mod_globals = temp_module.module.__dict__
82
        _run_code(code, mod_globals, init_globals,
Nick Coghlan's avatar
Nick Coghlan committed
83
                  mod_name, mod_fname, mod_loader, pkg_name)
84 85 86
    # Copy the globals of the temporary module, as they
    # may be cleared when the temporary module goes away
    return mod_globals.copy()
87 88


89 90
# This helper is needed due to a missing component in the PEP 302
# loader protocol (specifically, "get_filename" is non-standard)
91 92
# Since we can't introduce new features in maintenance releases,
# support was added to zipimporter under the name '_get_filename'
93
def _get_filename(loader, mod_name):
94 95 96
    for attr in ("get_filename", "_get_filename"):
        meth = getattr(loader, attr, None)
        if meth is not None:
97
            return os.path.abspath(meth(mod_name))
98
    return None
99

100 101
# Helper to get the loader, code and filename for a module
def _get_module_details(mod_name):
102
    loader = get_loader(mod_name)
103
    if loader is None:
104 105
        raise ImportError("No module named %s" % mod_name)
    if loader.is_package(mod_name):
106
        if mod_name == "__main__" or mod_name.endswith(".__main__"):
Benjamin Peterson's avatar
Benjamin Peterson committed
107
            raise ImportError("Cannot use package as __main__ module")
108 109 110 111 112 113
        try:
            pkg_main_name = mod_name + ".__main__"
            return _get_module_details(pkg_main_name)
        except ImportError as e:
            raise ImportError(("%s; %r is a package and cannot " +
                               "be directly executed") %(e, mod_name))
114 115
    code = loader.get_code(mod_name)
    if code is None:
116
        raise ImportError("No code object available for %s" % mod_name)
117
    filename = _get_filename(loader, mod_name)
118
    return mod_name, loader, code, filename
119

Benjamin Peterson's avatar
Benjamin Peterson committed
120 121 122 123
# XXX ncoghlan: Should this be documented and made public?
# (Current thoughts: don't repeat the mistake that lead to its
# creation when run_module() no longer met the needs of
# mainmodule.c, but couldn't be changed because it was public)
Nick Coghlan's avatar
Nick Coghlan committed
124
def _run_module_as_main(mod_name, alter_argv=True):
125 126
    """Runs the designated module in the __main__ namespace

Benjamin Peterson's avatar
Benjamin Peterson committed
127 128
       Note that the executed module will have full access to the
       __main__ namespace. If this is not desirable, the run_module()
129
       function should be used to run the module code in a fresh namespace.
Benjamin Peterson's avatar
Benjamin Peterson committed
130 131 132

       At the very least, these variables in __main__ will be overwritten:
           __name__
133
           __file__
Barry Warsaw's avatar
Barry Warsaw committed
134
           __cached__
135
           __loader__
Benjamin Peterson's avatar
Benjamin Peterson committed
136
           __package__
137
    """
Christian Heimes's avatar
Christian Heimes committed
138
    try:
Nick Coghlan's avatar
Nick Coghlan committed
139 140 141 142
        if alter_argv or mod_name != "__main__": # i.e. -m switch
            mod_name, loader, code, fname = _get_module_details(mod_name)
        else:          # i.e. directory or zipfile execution
            mod_name, loader, code, fname = _get_main_module_details()
Christian Heimes's avatar
Christian Heimes committed
143
    except ImportError as exc:
Benjamin Peterson's avatar
Benjamin Peterson committed
144 145 146 147 148 149 150 151
        # Try to provide a good error message
        # for directories, zip files and the -m switch
        if alter_argv:
            # For -m switch, just display the exception
            info = str(exc)
        else:
            # For directories/zipfiles, let the user
            # know what the code was looking for
152
            info = "can't find '__main__' module in %r" % sys.argv[0]
Benjamin Peterson's avatar
Benjamin Peterson committed
153
        msg = "%s: %s" % (sys.executable, info)
Christian Heimes's avatar
Christian Heimes committed
154
        sys.exit(msg)
155
    pkg_name = mod_name.rpartition('.')[0]
156
    main_globals = sys.modules["__main__"].__dict__
Nick Coghlan's avatar
Nick Coghlan committed
157
    if alter_argv:
158 159
        sys.argv[0] = fname
    return _run_code(code, main_globals, None,
160
                     "__main__", fname, loader, pkg_name)
161 162 163 164 165 166 167

def run_module(mod_name, init_globals=None,
               run_name=None, alter_sys=False):
    """Execute a module's code without importing it

       Returns the resulting top level namespace dictionary
    """
168
    mod_name, loader, code, fname = _get_module_details(mod_name)
169 170
    if run_name is None:
        run_name = mod_name
171
    pkg_name = mod_name.rpartition('.')[0]
172 173
    if alter_sys:
        return _run_module_code(code, init_globals, run_name,
174
                                fname, loader, pkg_name)
175 176
    else:
        # Leave the sys module alone
177 178
        return _run_code(code, {}, init_globals, run_name,
                         fname, loader, pkg_name)
179

Benjamin Peterson's avatar
Benjamin Peterson committed
180 181 182
def _get_main_module_details():
    # Helper that gives a nicer error message when attempting to
    # execute a zipfile or directory by invoking __main__.py
183 184
    # Also moves the standard __main__ out of the way so that the
    # preexisting __loader__ entry doesn't cause issues
Benjamin Peterson's avatar
Benjamin Peterson committed
185
    main_name = "__main__"
186 187
    saved_main = sys.modules[main_name]
    del sys.modules[main_name]
Benjamin Peterson's avatar
Benjamin Peterson committed
188 189 190 191 192
    try:
        return _get_module_details(main_name)
    except ImportError as exc:
        if main_name in str(exc):
            raise ImportError("can't find %r module in %r" %
193
                              (main_name, sys.path[0])) from exc
Benjamin Peterson's avatar
Benjamin Peterson committed
194
        raise
195 196
    finally:
        sys.modules[main_name] = saved_main
Benjamin Peterson's avatar
Benjamin Peterson committed
197

198

199
def _get_code_from_file(run_name, fname):
Nick Coghlan's avatar
Nick Coghlan committed
200 201 202 203 204
    # Check for a compiled file first
    with open(fname, "rb") as f:
        code = read_code(f)
    if code is None:
        # That didn't work, so try it as normal source code
205
        with open(fname, "rb") as f:
Nick Coghlan's avatar
Nick Coghlan committed
206
            code = compile(f.read(), fname, 'exec')
207 208 209 210
            loader = importlib.machinery.SourceFileLoader(run_name, fname)
    else:
        loader = importlib.machinery.SourcelessFileLoader(run_name, fname)
    return code, loader
Nick Coghlan's avatar
Nick Coghlan committed
211 212 213 214 215 216 217 218 219 220 221 222 223

def run_path(path_name, init_globals=None, run_name=None):
    """Execute code located at the specified filesystem location

       Returns the resulting top level namespace dictionary

       The file path may refer directly to a Python script (i.e.
       one that could be directly executed with execfile) or else
       it may refer to a zipfile or directory containing a top
       level __main__.py script.
    """
    if run_name is None:
        run_name = "<run_path>"
224
    pkg_name = run_name.rpartition(".")[0]
225
    importer = get_importer(path_name)
226
    if isinstance(importer, (type(None), imp.NullImporter)):
Nick Coghlan's avatar
Nick Coghlan committed
227 228
        # Not a valid sys.path entry, so run the code directly
        # execfile() doesn't help as we want to allow compiled files
229
        code, mod_loader = _get_code_from_file(run_name, path_name)
230
        return _run_module_code(code, init_globals, run_name, path_name,
231
                                mod_loader, pkg_name)
Nick Coghlan's avatar
Nick Coghlan committed
232 233 234 235 236 237 238 239 240 241 242
    else:
        # Importer is defined for path, so add it to
        # the start of sys.path
        sys.path.insert(0, path_name)
        try:
            # Here's where things are a little different from the run_module
            # case. There, we only had to replace the module in sys while the
            # code was running and doing so was somewhat optional. Here, we
            # have no choice and we have to remove it even while we read the
            # code. If we don't do this, a __loader__ attribute in the
            # existing __main__ module may prevent location of the new module.
243
            mod_name, loader, code, fname = _get_main_module_details()
Nick Coghlan's avatar
Nick Coghlan committed
244 245 246 247
            with _TempModule(run_name) as temp_module, \
                 _ModifiedArgv0(path_name):
                mod_globals = temp_module.module.__dict__
                return _run_code(code, mod_globals, init_globals,
248
                                    run_name, fname, loader, pkg_name).copy()
Nick Coghlan's avatar
Nick Coghlan committed
249 250 251 252 253 254 255
        finally:
            try:
                sys.path.remove(path_name)
            except ValueError:
                pass


256 257 258
if __name__ == "__main__":
    # Run the module specified as the next command line argument
    if len(sys.argv) < 2:
259
        print("No module specified for execution", file=sys.stderr)
260 261
    else:
        del sys.argv[0] # Make the requested module sys.argv[0]
262
        _run_module_as_main(sys.argv[0])