183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh#
283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh# Module for starting a process object using os.fork() or CreateProcess()
383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh#
483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh# multiprocessing/forking.py
583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh#
683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh# Copyright (c) 2006-2008, R Oudkerk
783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh# All rights reserved.
883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh#
983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh# Redistribution and use in source and binary forms, with or without
1083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh# modification, are permitted provided that the following conditions
1183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh# are met:
1283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh#
1383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh# 1. Redistributions of source code must retain the above copyright
1483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh#    notice, this list of conditions and the following disclaimer.
1583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh# 2. Redistributions in binary form must reproduce the above copyright
1683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh#    notice, this list of conditions and the following disclaimer in the
1783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh#    documentation and/or other materials provided with the distribution.
1883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh# 3. Neither the name of author nor the names of any contributors may be
1983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh#    used to endorse or promote products derived from this software
2083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh#    without specific prior written permission.
2183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh#
2283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh# THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS "AS IS" AND
2383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
2483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
2583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh# ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
2683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
2783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
2883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh# OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
2983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh# HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
3083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
3183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh# OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
3283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh# SUCH DAMAGE.
3383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh#
3483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
3583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehimport os
3683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehimport sys
3783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehimport signal
3883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehimport errno
3983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
4083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehfrom multiprocessing import util, process
4183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
4283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh__all__ = ['Popen', 'assert_spawning', 'exit', 'duplicate', 'close', 'ForkingPickler']
4383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
4483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh#
4583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh# Check that the current thread is spawning a child process
4683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh#
4783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
4883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehdef assert_spawning(self):
4983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    if not Popen.thread_is_spawning():
5083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        raise RuntimeError(
5183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            '%s objects should only be shared between processes'
5283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            ' through inheritance' % type(self).__name__
5383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            )
5483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
5583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh#
5683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh# Try making some callable types picklable
5783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh#
5883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
5983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehfrom pickle import Pickler
6083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehclass ForkingPickler(Pickler):
6183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    dispatch = Pickler.dispatch.copy()
6283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
6383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    @classmethod
6483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    def register(cls, type, reduce):
6583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        def dispatcher(self, obj):
6683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            rv = reduce(obj)
6783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            self.save_reduce(obj=obj, *rv)
6883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        cls.dispatch[type] = dispatcher
6983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
7083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehdef _reduce_method(m):
7183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    if m.im_self is None:
7283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        return getattr, (m.im_class, m.im_func.func_name)
7383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    else:
7483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        return getattr, (m.im_self, m.im_func.func_name)
7583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew HsiehForkingPickler.register(type(ForkingPickler.save), _reduce_method)
7683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
7783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehdef _reduce_method_descriptor(m):
7883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    return getattr, (m.__objclass__, m.__name__)
7983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew HsiehForkingPickler.register(type(list.append), _reduce_method_descriptor)
8083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew HsiehForkingPickler.register(type(int.__add__), _reduce_method_descriptor)
8183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
8283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh#def _reduce_builtin_function_or_method(m):
8383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh#    return getattr, (m.__self__, m.__name__)
8483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh#ForkingPickler.register(type(list().append), _reduce_builtin_function_or_method)
8583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh#ForkingPickler.register(type(int().__add__), _reduce_builtin_function_or_method)
8683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
8783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehtry:
8883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    from functools import partial
8983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehexcept ImportError:
9083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    pass
9183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehelse:
9283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    def _reduce_partial(p):
9383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        return _rebuild_partial, (p.func, p.args, p.keywords or {})
9483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    def _rebuild_partial(func, args, keywords):
9583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        return partial(func, *args, **keywords)
9683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    ForkingPickler.register(partial, _reduce_partial)
9783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
9883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh#
9983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh# Unix
10083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh#
10183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
10283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehif sys.platform != 'win32':
10383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    import time
10483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
10583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    exit = os._exit
10683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    duplicate = os.dup
10783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    close = os.close
10883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
10983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    #
11083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    # We define a Popen class similar to the one from subprocess, but
11183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    # whose constructor takes a process object as its argument.
11283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    #
11383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
11483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    class Popen(object):
11583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
11683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        def __init__(self, process_obj):
11783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            sys.stdout.flush()
11883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            sys.stderr.flush()
11983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            self.returncode = None
12083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
12183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            self.pid = os.fork()
12283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            if self.pid == 0:
12383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                if 'random' in sys.modules:
12483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                    import random
12583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                    random.seed()
12683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                code = process_obj._bootstrap()
12783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                sys.stdout.flush()
12883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                sys.stderr.flush()
12983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                os._exit(code)
13083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
13183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        def poll(self, flag=os.WNOHANG):
13283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            if self.returncode is None:
13383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                while True:
13483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                    try:
13583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                        pid, sts = os.waitpid(self.pid, flag)
13683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                    except os.error as e:
13783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                        if e.errno == errno.EINTR:
13883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                            continue
13983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                        # Child process not yet created. See #1731717
14083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                        # e.errno == errno.ECHILD == 10
14183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                        return None
14283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                    else:
14383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                        break
14483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                if pid == self.pid:
14583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                    if os.WIFSIGNALED(sts):
14683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                        self.returncode = -os.WTERMSIG(sts)
14783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                    else:
14883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                        assert os.WIFEXITED(sts)
14983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                        self.returncode = os.WEXITSTATUS(sts)
15083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            return self.returncode
15183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
15283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        def wait(self, timeout=None):
15383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            if timeout is None:
15483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                return self.poll(0)
15583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            deadline = time.time() + timeout
15683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            delay = 0.0005
15783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            while 1:
15883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                res = self.poll()
15983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                if res is not None:
16083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                    break
16183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                remaining = deadline - time.time()
16283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                if remaining <= 0:
16383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                    break
16483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                delay = min(delay * 2, remaining, 0.05)
16583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                time.sleep(delay)
16683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            return res
16783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
16883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        def terminate(self):
16983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            if self.returncode is None:
17083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                try:
17183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                    os.kill(self.pid, signal.SIGTERM)
17283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                except OSError, e:
17383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                    if self.wait(timeout=0.1) is None:
17483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                        raise
17583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
17683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        @staticmethod
17783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        def thread_is_spawning():
17883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            return False
17983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
18083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh#
18183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh# Windows
18283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh#
18383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
18483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehelse:
18583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    import thread
18683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    import msvcrt
18783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    import _subprocess
18883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    import time
18983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
19083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    from _multiprocessing import win32, Connection, PipeConnection
19183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    from .util import Finalize
19283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
19383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    #try:
19483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    #    from cPickle import dump, load, HIGHEST_PROTOCOL
19583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    #except ImportError:
19683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    from pickle import load, HIGHEST_PROTOCOL
19783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
19883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    def dump(obj, file, protocol=None):
19983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        ForkingPickler(file, protocol).dump(obj)
20083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
20183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    #
20283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    #
20383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    #
20483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
20583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    TERMINATE = 0x10000
20683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    WINEXE = (sys.platform == 'win32' and getattr(sys, 'frozen', False))
20783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    WINSERVICE = sys.executable.lower().endswith("pythonservice.exe")
20883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
20983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    exit = win32.ExitProcess
21083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    close = win32.CloseHandle
21183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
21283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    #
21383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    # _python_exe is the assumed path to the python executable.
21483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    # People embedding Python want to modify it.
21583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    #
21683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
21783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    if WINSERVICE:
21883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        _python_exe = os.path.join(sys.exec_prefix, 'python.exe')
21983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    else:
22083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        _python_exe = sys.executable
22183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
22283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    def set_executable(exe):
22383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        global _python_exe
22483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        _python_exe = exe
22583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
22683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    #
22783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    #
22883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    #
22983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
23083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    def duplicate(handle, target_process=None, inheritable=False):
23183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        if target_process is None:
23283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            target_process = _subprocess.GetCurrentProcess()
23383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        return _subprocess.DuplicateHandle(
23483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            _subprocess.GetCurrentProcess(), handle, target_process,
23583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            0, inheritable, _subprocess.DUPLICATE_SAME_ACCESS
23683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            ).Detach()
23783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
23883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    #
23983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    # We define a Popen class similar to the one from subprocess, but
24083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    # whose constructor takes a process object as its argument.
24183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    #
24283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
24383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    class Popen(object):
24483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        '''
24583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        Start a subprocess to run the code of a process object
24683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        '''
24783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        _tls = thread._local()
24883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
24983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        def __init__(self, process_obj):
25083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            # create pipe for communication with child
25183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            rfd, wfd = os.pipe()
25283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
25383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            # get handle for read end of the pipe and make it inheritable
25483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            rhandle = duplicate(msvcrt.get_osfhandle(rfd), inheritable=True)
25583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            os.close(rfd)
25683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
25783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            # start process
25883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            cmd = get_command_line() + [rhandle]
25983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            cmd = ' '.join('"%s"' % x for x in cmd)
26083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            hp, ht, pid, tid = _subprocess.CreateProcess(
26183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                _python_exe, cmd, None, None, 1, 0, None, None, None
26283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                )
26383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            ht.Close()
26483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            close(rhandle)
26583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
26683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            # set attributes of self
26783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            self.pid = pid
26883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            self.returncode = None
26983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            self._handle = hp
27083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
27183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            # send information to child
27283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            prep_data = get_preparation_data(process_obj._name)
27383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            to_child = os.fdopen(wfd, 'wb')
27483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            Popen._tls.process_handle = int(hp)
27583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            try:
27683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                dump(prep_data, to_child, HIGHEST_PROTOCOL)
27783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                dump(process_obj, to_child, HIGHEST_PROTOCOL)
27883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            finally:
27983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                del Popen._tls.process_handle
28083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                to_child.close()
28183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
28283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        @staticmethod
28383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        def thread_is_spawning():
28483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            return getattr(Popen._tls, 'process_handle', None) is not None
28583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
28683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        @staticmethod
28783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        def duplicate_for_child(handle):
28883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            return duplicate(handle, Popen._tls.process_handle)
28983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
29083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        def wait(self, timeout=None):
29183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            if self.returncode is None:
29283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                if timeout is None:
29383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                    msecs = _subprocess.INFINITE
29483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                else:
29583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                    msecs = max(0, int(timeout * 1000 + 0.5))
29683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
29783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                res = _subprocess.WaitForSingleObject(int(self._handle), msecs)
29883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                if res == _subprocess.WAIT_OBJECT_0:
29983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                    code = _subprocess.GetExitCodeProcess(self._handle)
30083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                    if code == TERMINATE:
30183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                        code = -signal.SIGTERM
30283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                    self.returncode = code
30383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
30483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            return self.returncode
30583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
30683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        def poll(self):
30783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            return self.wait(timeout=0)
30883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
30983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        def terminate(self):
31083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            if self.returncode is None:
31183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                try:
31283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                    _subprocess.TerminateProcess(int(self._handle), TERMINATE)
31383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                except WindowsError:
31483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                    if self.wait(timeout=0.1) is None:
31583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                        raise
31683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
31783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    #
31883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    #
31983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    #
32083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
32183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    def is_forking(argv):
32283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        '''
32383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        Return whether commandline indicates we are forking
32483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        '''
32583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        if len(argv) >= 2 and argv[1] == '--multiprocessing-fork':
32683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            assert len(argv) == 3
32783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            return True
32883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        else:
32983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            return False
33083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
33183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
33283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    def freeze_support():
33383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        '''
33483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        Run code for process object if this in not the main process
33583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        '''
33683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        if is_forking(sys.argv):
33783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            main()
33883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            sys.exit()
33983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
34083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
34183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    def get_command_line():
34283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        '''
34383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        Returns prefix of command line used for spawning a child process
34483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        '''
34583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        if getattr(process.current_process(), '_inheriting', False):
34683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            raise RuntimeError('''
34783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            Attempt to start a new process before the current process
34883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            has finished its bootstrapping phase.
34983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
35083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            This probably means that you are on Windows and you have
35183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            forgotten to use the proper idiom in the main module:
35283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
35383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                if __name__ == '__main__':
35483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                    freeze_support()
35583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                    ...
35683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
35783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            The "freeze_support()" line can be omitted if the program
35883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            is not going to be frozen to produce a Windows executable.''')
35983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
36083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        if getattr(sys, 'frozen', False):
36183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            return [sys.executable, '--multiprocessing-fork']
36283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        else:
36383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            prog = 'from multiprocessing.forking import main; main()'
36483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            opts = util._args_from_interpreter_flags()
36583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            return [_python_exe] + opts + ['-c', prog, '--multiprocessing-fork']
36683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
36783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
36883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    def main():
36983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        '''
37083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        Run code specifed by data received over pipe
37183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        '''
37283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        assert is_forking(sys.argv)
37383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
37483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        handle = int(sys.argv[-1])
37583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        fd = msvcrt.open_osfhandle(handle, os.O_RDONLY)
37683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        from_parent = os.fdopen(fd, 'rb')
37783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
37883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        process.current_process()._inheriting = True
37983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        preparation_data = load(from_parent)
38083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        prepare(preparation_data)
38183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        self = load(from_parent)
38283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        process.current_process()._inheriting = False
38383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
38483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        from_parent.close()
38583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
38683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        exitcode = self._bootstrap()
38783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        exit(exitcode)
38883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
38983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
39083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    def get_preparation_data(name):
39183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        '''
39283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        Return info about parent needed by child to unpickle process object
39383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        '''
39483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        from .util import _logger, _log_to_stderr
39583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
39683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        d = dict(
39783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            name=name,
39883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            sys_path=sys.path,
39983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            sys_argv=sys.argv,
40083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            log_to_stderr=_log_to_stderr,
40183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            orig_dir=process.ORIGINAL_DIR,
40283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            authkey=process.current_process().authkey,
40383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            )
40483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
40583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        if _logger is not None:
40683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            d['log_level'] = _logger.getEffectiveLevel()
40783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
40883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        if not WINEXE and not WINSERVICE:
40983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            main_path = getattr(sys.modules['__main__'], '__file__', None)
41083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            if not main_path and sys.argv[0] not in ('', '-c'):
41183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                main_path = sys.argv[0]
41283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            if main_path is not None:
41383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                if not os.path.isabs(main_path) and \
41483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                                          process.ORIGINAL_DIR is not None:
41583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                    main_path = os.path.join(process.ORIGINAL_DIR, main_path)
41683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                d['main_path'] = os.path.normpath(main_path)
41783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
41883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        return d
41983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
42083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    #
42183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    # Make (Pipe)Connection picklable
42283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    #
42383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
42483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    def reduce_connection(conn):
42583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        if not Popen.thread_is_spawning():
42683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            raise RuntimeError(
42783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                'By default %s objects can only be shared between processes\n'
42883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                'using inheritance' % type(conn).__name__
42983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                )
43083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        return type(conn), (Popen.duplicate_for_child(conn.fileno()),
43183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                            conn.readable, conn.writable)
43283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
43383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    ForkingPickler.register(Connection, reduce_connection)
43483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    ForkingPickler.register(PipeConnection, reduce_connection)
43583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
43683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh#
43783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh# Prepare current process
43883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh#
43983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
44083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehold_main_modules = []
44183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
44283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsiehdef prepare(data):
44383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    '''
44483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    Try to get current process ready to unpickle process object
44583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    '''
44683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    old_main_modules.append(sys.modules['__main__'])
44783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
44883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    if 'name' in data:
44983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        process.current_process().name = data['name']
45083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
45183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    if 'authkey' in data:
45283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        process.current_process()._authkey = data['authkey']
45383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
45483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    if 'log_to_stderr' in data and data['log_to_stderr']:
45583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        util.log_to_stderr()
45683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
45783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    if 'log_level' in data:
45883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        util.get_logger().setLevel(data['log_level'])
45983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
46083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    if 'sys_path' in data:
46183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        sys.path = data['sys_path']
46283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
46383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    if 'sys_argv' in data:
46483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        sys.argv = data['sys_argv']
46583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
46683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    if 'dir' in data:
46783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        os.chdir(data['dir'])
46883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
46983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    if 'orig_dir' in data:
47083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        process.ORIGINAL_DIR = data['orig_dir']
47183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
47283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh    if 'main_path' in data:
47383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        main_path = data['main_path']
47483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        main_name = os.path.splitext(os.path.basename(main_path))[0]
47583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        if main_name == '__init__':
47683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            main_name = os.path.basename(os.path.dirname(main_path))
47783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
47883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh        if main_name != 'ipython':
47983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            import imp
48083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
48183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            if main_path is None:
48283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                dirs = None
48383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            elif os.path.basename(main_path).startswith('__init__.py'):
48483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                dirs = [os.path.dirname(os.path.dirname(main_path))]
48583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            else:
48683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                dirs = [os.path.dirname(main_path)]
48783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
48883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            assert main_name not in sys.modules, main_name
48983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            file, path_name, etc = imp.find_module(main_name, dirs)
49083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            try:
49183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                # We would like to do "imp.load_module('__main__', ...)"
49283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                # here.  However, that would cause 'if __name__ ==
49383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                # "__main__"' clauses to be executed.
49483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                main_module = imp.load_module(
49583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                    '__parents_main__', file, path_name, etc
49683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                    )
49783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            finally:
49883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                if file:
49983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                    file.close()
50083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
50183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            sys.modules['__main__'] = main_module
50283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            main_module.__name__ = '__main__'
50383760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh
50483760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            # Try to make the potentially picklable objects in
50583760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            # sys.modules['__main__'] realize they are in the main
50683760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            # module -- somewhat ugly.
50783760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh            for obj in main_module.__dict__.values():
50883760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                try:
50983760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                    if obj.__module__ == '__parents_main__':
51083760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                        obj.__module__ = '__main__'
51183760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                except Exception:
51283760d213fb3bec7b4117d266fcfbf6fe2ba14abAndrew Hsieh                    pass
513