Lib/multiprocessing/reduction.py - platform/external/python/cpython3 - Gitiles

 #
 # Module which deals with pickling of objects.
 #
 # multiprocessing/reduction.py
 #
 # Copyright (c) 2006-2008, R Oudkerk
 # Licensed to PSF under a Contributor Agreement.
 #

 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')))

 #
 # Pickler subclass
 #

 class ForkingPickler(pickle.Pickler):
     '''Pickler subclass used by multiprocessing.'''
     _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):
         '''Register a reduce function for a type.'''
         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):
     '''Replacement for pickle.dump() using ForkingPickler.'''
     ForkingPickler(file, protocol).dump(obj)

 #
 # Platform specific definitions
 #

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

     def duplicate(handle, target_process=None, inheritable=False,
                   *, source_process=None):
         '''Duplicate a handle.  (target_process is a handle not a pid!)'''
         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):
         '''Steal a handle from process identified by source_pid.'''
         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):
         '''Send a handle over a local connection.'''
         dh = DupHandle(handle, _winapi.DUPLICATE_SAME_ACCESS, destination_pid)
         conn.send(dh)

     def recv_handle(conn):
         '''Receive a handle over a local connection.'''
         return conn.recv().detach()

     class DupHandle(object):
         '''Picklable wrapper for a handle.'''
         def __init__(self, handle, access, pid=None):
             if pid is None:
                 # We just duplicate the handle in the current process and
                 # let the receiving process steal the handle.
                 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):
             '''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.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:
     # Unix
     __all__ += ['DupFd', 'sendfds', 'recvfds']
     import array

     # On MacOSX we should acknowledge receipt of fds -- see Issue14669
     ACKNOWLEDGE = sys.platform == 'darwin'

     def sendfds(sock, fds):
         '''Send an array of fds over an AF_UNIX socket.'''
         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):
         '''Receive an array of fds over an AF_UNIX socket.'''
         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):
         '''Send a handle over a local connection.'''
         with socket.fromfd(conn.fileno(), socket.AF_UNIX, socket.SOCK_STREAM) as s:
             sendfds(s, [handle])

     def recv_handle(conn):
         '''Receive a handle over a local connection.'''
         with socket.fromfd(conn.fileno(), socket.AF_UNIX, socket.SOCK_STREAM) as s:
             return recvfds(s, 1)[0]

     def DupFd(fd):
         '''Return a wrapper for an 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')

 #
 # Try making some callable types picklable
 #

 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)

 #
 # Make sockets picklable
 #

 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):
     '''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
     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)
	#
	# Module which deals with pickling of objects.
	#
	# multiprocessing/reduction.py
	#
	# Copyright (c) 2006-2008, R Oudkerk
	# Licensed to PSF under a Contributor Agreement.
	#

	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')))

	#
	# Pickler subclass
	#

	class ForkingPickler(pickle.Pickler):
	'''Pickler subclass used by multiprocessing.'''
	_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):
	'''Register a reduce function for a type.'''
	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):
	'''Replacement for pickle.dump() using ForkingPickler.'''
	ForkingPickler(file, protocol).dump(obj)

	#
	# Platform specific definitions
	#

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

	def duplicate(handle, target_process=None, inheritable=False,
	*, source_process=None):
	'''Duplicate a handle. (target_process is a handle not a pid!)'''
	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):
	'''Steal a handle from process identified by source_pid.'''
	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):
	'''Send a handle over a local connection.'''
	dh = DupHandle(handle, _winapi.DUPLICATE_SAME_ACCESS, destination_pid)
	conn.send(dh)

	def recv_handle(conn):
	'''Receive a handle over a local connection.'''
	return conn.recv().detach()

	class DupHandle(object):
	'''Picklable wrapper for a handle.'''
	def __init__(self, handle, access, pid=None):
	if pid is None:
	# We just duplicate the handle in the current process and
	# let the receiving process steal the handle.
	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):
	'''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.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:
	# Unix
	__all__ += ['DupFd', 'sendfds', 'recvfds']
	import array

	# On MacOSX we should acknowledge receipt of fds -- see Issue14669
	ACKNOWLEDGE = sys.platform == 'darwin'

	def sendfds(sock, fds):
	'''Send an array of fds over an AF_UNIX socket.'''
	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):
	'''Receive an array of fds over an AF_UNIX socket.'''
	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):
	'''Send a handle over a local connection.'''
	with socket.fromfd(conn.fileno(), socket.AF_UNIX, socket.SOCK_STREAM) as s:
	sendfds(s, [handle])

	def recv_handle(conn):
	'''Receive a handle over a local connection.'''
	with socket.fromfd(conn.fileno(), socket.AF_UNIX, socket.SOCK_STREAM) as s:
	return recvfds(s, 1)[0]

	def DupFd(fd):
	'''Return a wrapper for an 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')

	#
	# Try making some callable types picklable
	#

	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)

	#
	# Make sockets picklable
	#

	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):
	'''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
	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)