1#
2# Code used to start processes when using the spawn or forkserver
3# start methods.
4#
5# multiprocessing/spawn.py
6#
7# Copyright (c) 2006-2008, R Oudkerk
8# Licensed to PSF under a Contributor Agreement.
9#
10
11import os
12import sys
13import runpy
14import types
15
16from . import get_start_method, set_start_method
17from . import process
18from .context import reduction
19from . import util
20
21__all__ = ['_main', 'freeze_support', 'set_executable', 'get_executable',
22           'get_preparation_data', 'get_command_line', 'import_main_path']
23
24#
25# _python_exe is the assumed path to the python executable.
26# People embedding Python want to modify it.
27#
28
29if sys.platform != 'win32':
30    WINEXE = False
31    WINSERVICE = False
32else:
33    WINEXE = getattr(sys, 'frozen', False)
34    WINSERVICE = sys.executable.lower().endswith("pythonservice.exe")
35
36if WINSERVICE:
37    _python_exe = os.path.join(sys.exec_prefix, 'python.exe')
38else:
39    _python_exe = sys.executable
40
41def set_executable(exe):
42    global _python_exe
43    _python_exe = exe
44
45def get_executable():
46    return _python_exe
47
48#
49#
50#
51
52def is_forking(argv):
53    '''
54    Return whether commandline indicates we are forking
55    '''
56    if len(argv) >= 2 and argv[1] == '--multiprocessing-fork':
57        return True
58    else:
59        return False
60
61
62def freeze_support():
63    '''
64    Run code for process object if this in not the main process
65    '''
66    if is_forking(sys.argv):
67        kwds = {}
68        for arg in sys.argv[2:]:
69            name, value = arg.split('=')
70            if value == 'None':
71                kwds[name] = None
72            else:
73                kwds[name] = int(value)
74        spawn_main(**kwds)
75        sys.exit()
76
77
78def get_command_line(**kwds):
79    '''
80    Returns prefix of command line used for spawning a child process
81    '''
82    if getattr(sys, 'frozen', False):
83        return ([sys.executable, '--multiprocessing-fork'] +
84                ['%s=%r' % item for item in kwds.items()])
85    else:
86        prog = 'from multiprocess.spawn import spawn_main; spawn_main(%s)'
87        prog %= ', '.join('%s=%r' % item for item in kwds.items())
88        opts = util._args_from_interpreter_flags()
89        return [_python_exe] + opts + ['-c', prog, '--multiprocessing-fork']
90
91
92def spawn_main(pipe_handle, parent_pid=None, tracker_fd=None):
93    '''
94    Run code specified by data received over pipe
95    '''
96    assert is_forking(sys.argv), "Not forking"
97    if sys.platform == 'win32':
98        import msvcrt
99        new_handle = reduction.steal_handle(parent_pid, pipe_handle)
100        fd = msvcrt.open_osfhandle(new_handle, os.O_RDONLY)
101    else:
102        from . import semaphore_tracker
103        semaphore_tracker._semaphore_tracker._fd = tracker_fd
104        fd = pipe_handle
105    exitcode = _main(fd)
106    sys.exit(exitcode)
107
108
109def _main(fd):
110    with os.fdopen(fd, 'rb', closefd=True) as from_parent:
111        process.current_process()._inheriting = True
112        try:
113            preparation_data = reduction.pickle.load(from_parent)
114            prepare(preparation_data)
115            self = reduction.pickle.load(from_parent)
116        finally:
117            del process.current_process()._inheriting
118    return self._bootstrap()
119
120
121def _check_not_importing_main():
122    if getattr(process.current_process(), '_inheriting', False):
123        raise RuntimeError('''
124        An attempt has been made to start a new process before the
125        current process has finished its bootstrapping phase.
126
127        This probably means that you are not using fork to start your
128        child processes and you have forgotten to use the proper idiom
129        in the main module:
130
131            if __name__ == '__main__':
132                freeze_support()
133                ...
134
135        The "freeze_support()" line can be omitted if the program
136        is not going to be frozen to produce an executable.''')
137
138
139def get_preparation_data(name):
140    '''
141    Return info about parent needed by child to unpickle process object
142    '''
143    _check_not_importing_main()
144    d = dict(
145        log_to_stderr=util._log_to_stderr,
146        authkey=process.current_process().authkey,
147        )
148
149    if util._logger is not None:
150        d['log_level'] = util._logger.getEffectiveLevel()
151
152    sys_path=sys.path.copy()
153    try:
154        i = sys_path.index('')
155    except ValueError:
156        pass
157    else:
158        sys_path[i] = process.ORIGINAL_DIR
159
160    d.update(
161        name=name,
162        sys_path=sys_path,
163        sys_argv=sys.argv,
164        orig_dir=process.ORIGINAL_DIR,
165        dir=os.getcwd(),
166        start_method=get_start_method(),
167        )
168
169    # Figure out whether to initialise main in the subprocess as a module
170    # or through direct execution (or to leave it alone entirely)
171    main_module = sys.modules['__main__']
172    main_mod_name = getattr(main_module.__spec__, "name", None)
173    if main_mod_name is not None:
174        d['init_main_from_name'] = main_mod_name
175    elif sys.platform != 'win32' or (not WINEXE and not WINSERVICE):
176        main_path = getattr(main_module, '__file__', None)
177        if main_path is not None:
178            if (not os.path.isabs(main_path) and
179                        process.ORIGINAL_DIR is not None):
180                main_path = os.path.join(process.ORIGINAL_DIR, main_path)
181            d['init_main_from_path'] = os.path.normpath(main_path)
182
183    return d
184
185#
186# Prepare current process
187#
188
189old_main_modules = []
190
191def prepare(data):
192    '''
193    Try to get current process ready to unpickle process object
194    '''
195    if 'name' in data:
196        process.current_process().name = data['name']
197
198    if 'authkey' in data:
199        process.current_process().authkey = data['authkey']
200
201    if 'log_to_stderr' in data and data['log_to_stderr']:
202        util.log_to_stderr()
203
204    if 'log_level' in data:
205        util.get_logger().setLevel(data['log_level'])
206
207    if 'sys_path' in data:
208        sys.path = data['sys_path']
209
210    if 'sys_argv' in data:
211        sys.argv = data['sys_argv']
212
213    if 'dir' in data:
214        os.chdir(data['dir'])
215
216    if 'orig_dir' in data:
217        process.ORIGINAL_DIR = data['orig_dir']
218
219    if 'start_method' in data:
220        set_start_method(data['start_method'], force=True)
221
222    if 'init_main_from_name' in data:
223        _fixup_main_from_name(data['init_main_from_name'])
224    elif 'init_main_from_path' in data:
225        _fixup_main_from_path(data['init_main_from_path'])
226
227# Multiprocessing module helpers to fix up the main module in
228# spawned subprocesses
229def _fixup_main_from_name(mod_name):
230    # __main__.py files for packages, directories, zip archives, etc, run
231    # their "main only" code unconditionally, so we don't even try to
232    # populate anything in __main__, nor do we make any changes to
233    # __main__ attributes
234    current_main = sys.modules['__main__']
235    if mod_name == "__main__" or mod_name.endswith(".__main__"):
236        return
237
238    # If this process was forked, __main__ may already be populated
239    if getattr(current_main.__spec__, "name", None) == mod_name:
240        return
241
242    # Otherwise, __main__ may contain some non-main code where we need to
243    # support unpickling it properly. We rerun it as __mp_main__ and make
244    # the normal __main__ an alias to that
245    old_main_modules.append(current_main)
246    main_module = types.ModuleType("__mp_main__")
247    main_content = runpy.run_module(mod_name,
248                                    run_name="__mp_main__",
249                                    alter_sys=True)
250    main_module.__dict__.update(main_content)
251    sys.modules['__main__'] = sys.modules['__mp_main__'] = main_module
252
253
254def _fixup_main_from_path(main_path):
255    # If this process was forked, __main__ may already be populated
256    current_main = sys.modules['__main__']
257
258    # Unfortunately, the main ipython launch script historically had no
259    # "if __name__ == '__main__'" guard, so we work around that
260    # by treating it like a __main__.py file
261    # See https://github.com/ipython/ipython/issues/4698
262    main_name = os.path.splitext(os.path.basename(main_path))[0]
263    if main_name == 'ipython':
264        return
265
266    # Otherwise, if __file__ already has the setting we expect,
267    # there's nothing more to do
268    if getattr(current_main, '__file__', None) == main_path:
269        return
270
271    # If the parent process has sent a path through rather than a module
272    # name we assume it is an executable script that may contain
273    # non-main code that needs to be executed
274    old_main_modules.append(current_main)
275    main_module = types.ModuleType("__mp_main__")
276    main_content = runpy.run_path(main_path,
277                                  run_name="__mp_main__")
278    main_module.__dict__.update(main_content)
279    sys.modules['__main__'] = sys.modules['__mp_main__'] = main_module
280
281
282def import_main_path(main_path):
283    '''
284    Set sys.modules['__main__'] to module at main_path
285    '''
286    _fixup_main_from_path(main_path)
287