module which deals with pickling of objects multiprocessingreduction py c 20062008 r oudkerk licensed to psf under a contributor agreement pickler subclass pickler subclass used by multiprocessing extrareducers copyregdispatchtable copyreg dispatchtable def initself args super initargs self dispatchtable self copyregdispatchtable copy self dispatchtable updateself extrareducers classmethod def registercls type reduce replacement for pickle dump using forkingpickler forkingpicklerfile protocol dumpobj platform specific definitions if sys platform win32 windows all duphandle duplicate stealhandle import winapi def duplicatehandle targetprocessnone inheritablefalse sourceprocessnone steal a handle from process identified by sourcepid sourceprocesshandle winapi openprocess winapi processduphandle false sourcepid try return winapi duplicatehandle sourceprocesshandle handle winapi getcurrentprocess 0 false winapi duplicatesameaccess winapi duplicateclosesource finally winapi closehandlesourceprocesshandle def sendhandleconn handle destinationpid receive a handle over a local connection return conn recv detach class duphandleobject we just duplicate the handle in the current process and let the receiving process steal the handle get the handle this should only be called once retrieve handle from process which currently owns it if self pid os getpid the handle has already been duplicated for this process return self handle we must steal the handle from the process whose pid is self pid proc winapi openprocesswinapi processduphandle false self pid try return winapi duplicatehandle proc self handle winapi getcurrentprocess self access false winapi duplicateclosesource finally winapi closehandleproc else unix all dupfd sendfds recvfds import array on macosx we should acknowledge receipt of fds see issue14669 acknowledge sys platform darwin def sendfdssock fds receive an array of fds over an afunix socket a array array i bytessize a itemsize size msg ancdata flags addr sock recvmsg1 socket cmsgspacebytessize if not msg and not ancdata raise eoferror try if acknowledge sock sendb a if lenancdata 1 raise runtimeerror received d items of ancdata lenancdata cmsglevel cmsgtype cmsgdata ancdata0 if cmsglevel socket solsocket and cmsgtype socket scmrights if lencmsgdata a itemsize 0 raise valueerror a frombytescmsgdata if lena 256 msg0 raise assertionerror len is 0 n but msg0 is 1 r format lena msg0 return lista except valueerror indexerror pass raise runtimeerror invalid data received def sendhandleconn handle destinationpid receive a handle over a local connection with socket fromfdconn fileno socket afunix socket sockstream as s return recvfdss 10 def dupfdfd try making some callable types picklable make sockets picklable abstract base class for use in implementing a reduction class suitable for use in replacing the standard reduction mechanism used in multiprocessing forkingpickler forkingpickler register register dump dump sendhandle sendhandle recvhandle recvhandle if sys platform win32 stealhandle stealhandle duplicate duplicate duphandle duphandle else sendfds sendfds recvfds recvfds dupfd dupfd reducemethod reducemethod reducemethoddescriptor reducemethoddescriptor rebuildpartial rebuildpartial reducesocket reducesocket rebuildsocket rebuildsocket def initself args registertypec f reducemethod registertypelist append reducemethoddescriptor registertypeint add reducemethoddescriptor registerfunctools partial reducepartial registersocket socket reducesocket module which deals with pickling of objects multiprocessing reduction py c 2006 2008 r oudkerk licensed to psf under a contributor agreement pickler subclass pickler subclass used by multiprocessing register a reduce function for a type replacement for pickle dump using forkingpickler platform specific definitions windows duplicate a handle target_process is a handle not a pid steal a handle from process identified by source_pid send a handle over a local connection receive a handle over a local connection picklable wrapper for a handle we just duplicate the handle in the current process and let the receiving process steal the handle get the handle this should only be called once retrieve handle from process which currently owns it the handle has already been duplicated for this process we must steal the handle from the process whose pid is self _pid unix on macosx we should acknowledge receipt of fds see issue14669 send an array of fds over an af_unix socket receive an array of fds over an af_unix socket send a handle over a local connection receive a handle over a local connection return a wrapper for an fd try making some callable types picklable make sockets picklable abstract base class for use in implementing a reduction class suitable for use in replacing the standard reduction mechanism used in multiprocessing
from abc import ABCMeta
import copyreg
import functools
import io
import os
import pickle
import socket
import sys

from . import context

__all__ = ['send_handle', 'recv_handle', 'ForkingPickler', 'register', 'dump']


HAVE_SEND_HANDLE = (sys.platform == 'win32' or
                    (hasattr(socket, 'CMSG_LEN') and
                     hasattr(socket, 'SCM_RIGHTS') and
                     hasattr(socket.socket, 'sendmsg')))





class ForkingPickler(pickle.Pickler):
    
    _extra_reducers = {}
    _copyreg_dispatch_table = copyreg.dispatch_table

    def __init__(self, *args):
        super().__init__(*args)
        self.dispatch_table = self._copyreg_dispatch_table.copy()
        self.dispatch_table.update(self._extra_reducers)

    @classmethod
    def register(cls, type, reduce):
        
        cls._extra_reducers[type] = reduce

    @classmethod
    def dumps(cls, obj, protocol=None):
        buf = io.BytesIO()
        cls(buf, protocol).dump(obj)
        return buf.getbuffer()

    loads = pickle.loads

register = ForkingPickler.register

def dump(obj, file, protocol=None):
    
    ForkingPickler(file, protocol).dump(obj)





if sys.platform == 'win32':
    
    __all__ += ['DupHandle', 'duplicate', 'steal_handle']
    import _winapi

    def duplicate(handle, target_process=None, inheritable=False,
                  *, source_process=None):
        
        current_process = _winapi.GetCurrentProcess()
        if source_process is None:
            source_process = current_process
        if target_process is None:
            target_process = current_process
        return _winapi.DuplicateHandle(
            source_process, handle, target_process,
            0, inheritable, _winapi.DUPLICATE_SAME_ACCESS)

    def steal_handle(source_pid, handle):
        
        source_process_handle = _winapi.OpenProcess(
            _winapi.PROCESS_DUP_HANDLE, False, source_pid)
        try:
            return _winapi.DuplicateHandle(
                source_process_handle, handle,
                _winapi.GetCurrentProcess(), 0, False,
                _winapi.DUPLICATE_SAME_ACCESS | _winapi.DUPLICATE_CLOSE_SOURCE)
        finally:
            _winapi.CloseHandle(source_process_handle)

    def send_handle(conn, handle, destination_pid):
        
        dh = DupHandle(handle, _winapi.DUPLICATE_SAME_ACCESS, destination_pid)
        conn.send(dh)

    def recv_handle(conn):
        
        return conn.recv().detach()

    class DupHandle(object):
        
        def __init__(self, handle, access, pid=None):
            if pid is None:
                
                
                pid = os.getpid()
            proc = _winapi.OpenProcess(_winapi.PROCESS_DUP_HANDLE, False, pid)
            try:
                self._handle = _winapi.DuplicateHandle(
                    _winapi.GetCurrentProcess(),
                    handle, proc, access, False, 0)
            finally:
                _winapi.CloseHandle(proc)
            self._access = access
            self._pid = pid

        def detach(self):
            
            
            if self._pid == os.getpid():
                
                return self._handle
            
            proc = _winapi.OpenProcess(_winapi.PROCESS_DUP_HANDLE, False,
                                       self._pid)
            try:
                return _winapi.DuplicateHandle(
                    proc, self._handle, _winapi.GetCurrentProcess(),
                    self._access, False, _winapi.DUPLICATE_CLOSE_SOURCE)
            finally:
                _winapi.CloseHandle(proc)

else:
    
    __all__ += ['DupFd', 'sendfds', 'recvfds']
    import array

    
    ACKNOWLEDGE = sys.platform == 'darwin'

    def sendfds(sock, fds):
        
        fds = array.array('i', fds)
        msg = bytes([len(fds) % 256])
        sock.sendmsg([msg], [(socket.SOL_SOCKET, socket.SCM_RIGHTS, fds)])
        if ACKNOWLEDGE and sock.recv(1) != b'A':
            raise RuntimeError('did not receive acknowledgement of fd')

    def recvfds(sock, size):
        
        a = array.array('i')
        bytes_size = a.itemsize * size
        msg, ancdata, flags, addr = sock.recvmsg(1, socket.CMSG_SPACE(bytes_size))
        if not msg and not ancdata:
            raise EOFError
        try:
            if ACKNOWLEDGE:
                sock.send(b'A')
            if len(ancdata) != 1:
                raise RuntimeError('received %d items of ancdata' %
                                   len(ancdata))
            cmsg_level, cmsg_type, cmsg_data = ancdata[0]
            if (cmsg_level == socket.SOL_SOCKET and
                cmsg_type == socket.SCM_RIGHTS):
                if len(cmsg_data) % a.itemsize != 0:
                    raise ValueError
                a.frombytes(cmsg_data)
                if len(a) % 256 != msg[0]:
                    raise AssertionError(
                        "Len is {0:n} but msg[0] is {1!r}".format(
                            len(a), msg[0]))
                return list(a)
        except (ValueError, IndexError):
            pass
        raise RuntimeError('Invalid data received')

    def send_handle(conn, handle, destination_pid):
        
        with socket.fromfd(conn.fileno(), socket.AF_UNIX, socket.SOCK_STREAM) as s:
            sendfds(s, [handle])

    def recv_handle(conn):
        
        with socket.fromfd(conn.fileno(), socket.AF_UNIX, socket.SOCK_STREAM) as s:
            return recvfds(s, 1)[0]

    def DupFd(fd):
        
        popen_obj = context.get_spawning_popen()
        if popen_obj is not None:
            return popen_obj.DupFd(popen_obj.duplicate_for_child(fd))
        elif HAVE_SEND_HANDLE:
            from . import resource_sharer
            return resource_sharer.DupFd(fd)
        else:
            raise ValueError('SCM_RIGHTS appears not to be available')





def _reduce_method(m):
    if m.__self__ is None:
        return getattr, (m.__class__, m.__func__.__name__)
    else:
        return getattr, (m.__self__, m.__func__.__name__)
class _C:
    def f(self):
        pass
register(type(_C().f), _reduce_method)


def _reduce_method_descriptor(m):
    return getattr, (m.__objclass__, m.__name__)
register(type(list.append), _reduce_method_descriptor)
register(type(int.__add__), _reduce_method_descriptor)


def _reduce_partial(p):
    return _rebuild_partial, (p.func, p.args, p.keywords or {})
def _rebuild_partial(func, args, keywords):
    return functools.partial(func, *args, **keywords)
register(functools.partial, _reduce_partial)





if sys.platform == 'win32':
    def _reduce_socket(s):
        from .resource_sharer import DupSocket
        return _rebuild_socket, (DupSocket(s),)
    def _rebuild_socket(ds):
        return ds.detach()
    register(socket.socket, _reduce_socket)

else:
    def _reduce_socket(s):
        df = DupFd(s.fileno())
        return _rebuild_socket, (df, s.family, s.type, s.proto)
    def _rebuild_socket(df, family, type, proto):
        fd = df.detach()
        return socket.socket(family, type, proto, fileno=fd)
    register(socket.socket, _reduce_socket)


class AbstractReducer(metaclass=ABCMeta):
    
    ForkingPickler = ForkingPickler
    register = register
    dump = dump
    send_handle = send_handle
    recv_handle = recv_handle

    if sys.platform == 'win32':
        steal_handle = steal_handle
        duplicate = duplicate
        DupHandle = DupHandle
    else:
        sendfds = sendfds
        recvfds = recvfds
        DupFd = DupFd

    _reduce_method = _reduce_method
    _reduce_method_descriptor = _reduce_method_descriptor
    _rebuild_partial = _rebuild_partial
    _reduce_socket = _reduce_socket
    _rebuild_socket = _rebuild_socket

    def __init__(self, *args):
        register(type(_C().f), _reduce_method)
        register(type(list.append), _reduce_method_descriptor)
        register(type(int.__add__), _reduce_method_descriptor)
        register(functools.partial, _reduce_partial)
        register(socket.socket, _reduce_socket)