Python code coverage for Lib/multiprocessing/reduction.py

#	count	content
1	n/a	#
2	n/a	# Module which deals with pickling of objects.
3	n/a	#
4	n/a	# multiprocessing/reduction.py
5	n/a	#
6	n/a	# Copyright (c) 2006-2008, R Oudkerk
7	n/a	# Licensed to PSF under a Contributor Agreement.
8	n/a	#
9	n/a
10	n/a	from abc import ABCMeta, abstractmethod
11	n/a	import copyreg
12	n/a	import functools
13	n/a	import io
14	n/a	import os
15	n/a	import pickle
16	n/a	import socket
17	n/a	import sys
18	n/a
19	n/a	from . import context
20	n/a
21	n/a	__all__ = ['send_handle', 'recv_handle', 'ForkingPickler', 'register', 'dump']
22	n/a
23	n/a
24	n/a	HAVE_SEND_HANDLE = (sys.platform == 'win32' or
25	n/a	(hasattr(socket, 'CMSG_LEN') and
26	n/a	hasattr(socket, 'SCM_RIGHTS') and
27	n/a	hasattr(socket.socket, 'sendmsg')))
28	n/a
29	n/a	#
30	n/a	# Pickler subclass
31	n/a	#
32	n/a
33	n/a	class ForkingPickler(pickle.Pickler):
34	n/a	'''Pickler subclass used by multiprocessing.'''
35	n/a	_extra_reducers = {}
36	n/a	_copyreg_dispatch_table = copyreg.dispatch_table
37	n/a
38	n/a	def __init__(self, *args):
39	n/a	super().__init__(*args)
40	n/a	self.dispatch_table = self._copyreg_dispatch_table.copy()
41	n/a	self.dispatch_table.update(self._extra_reducers)
42	n/a
43	n/a	@classmethod
44	n/a	def register(cls, type, reduce):
45	n/a	'''Register a reduce function for a type.'''
46	n/a	cls._extra_reducers[type] = reduce
47	n/a
48	n/a	@classmethod
49	n/a	def dumps(cls, obj, protocol=None):
50	n/a	buf = io.BytesIO()
51	n/a	cls(buf, protocol).dump(obj)
52	n/a	return buf.getbuffer()
53	n/a
54	n/a	loads = pickle.loads
55	n/a
56	n/a	register = ForkingPickler.register
57	n/a
58	n/a	def dump(obj, file, protocol=None):
59	n/a	'''Replacement for pickle.dump() using ForkingPickler.'''
60	n/a	ForkingPickler(file, protocol).dump(obj)
61	n/a
62	n/a	#
63	n/a	# Platform specific definitions
64	n/a	#
65	n/a
66	n/a	if sys.platform == 'win32':
67	n/a	# Windows
68	n/a	__all__ += ['DupHandle', 'duplicate', 'steal_handle']
69	n/a	import _winapi
70	n/a
71	n/a	def duplicate(handle, target_process=None, inheritable=False):
72	n/a	'''Duplicate a handle. (target_process is a handle not a pid!)'''
73	n/a	if target_process is None:
74	n/a	target_process = _winapi.GetCurrentProcess()
75	n/a	return _winapi.DuplicateHandle(
76	n/a	_winapi.GetCurrentProcess(), handle, target_process,
77	n/a	0, inheritable, _winapi.DUPLICATE_SAME_ACCESS)
78	n/a
79	n/a	def steal_handle(source_pid, handle):
80	n/a	'''Steal a handle from process identified by source_pid.'''
81	n/a	source_process_handle = _winapi.OpenProcess(
82	n/a	_winapi.PROCESS_DUP_HANDLE, False, source_pid)
83	n/a	try:
84	n/a	return _winapi.DuplicateHandle(
85	n/a	source_process_handle, handle,
86	n/a	_winapi.GetCurrentProcess(), 0, False,
87	n/a	_winapi.DUPLICATE_SAME_ACCESS \| _winapi.DUPLICATE_CLOSE_SOURCE)
88	n/a	finally:
89	n/a	_winapi.CloseHandle(source_process_handle)
90	n/a
91	n/a	def send_handle(conn, handle, destination_pid):
92	n/a	'''Send a handle over a local connection.'''
93	n/a	dh = DupHandle(handle, _winapi.DUPLICATE_SAME_ACCESS, destination_pid)
94	n/a	conn.send(dh)
95	n/a
96	n/a	def recv_handle(conn):
97	n/a	'''Receive a handle over a local connection.'''
98	n/a	return conn.recv().detach()
99	n/a
100	n/a	class DupHandle(object):
101	n/a	'''Picklable wrapper for a handle.'''
102	n/a	def __init__(self, handle, access, pid=None):
103	n/a	if pid is None:
104	n/a	# We just duplicate the handle in the current process and
105	n/a	# let the receiving process steal the handle.
106	n/a	pid = os.getpid()
107	n/a	proc = _winapi.OpenProcess(_winapi.PROCESS_DUP_HANDLE, False, pid)
108	n/a	try:
109	n/a	self._handle = _winapi.DuplicateHandle(
110	n/a	_winapi.GetCurrentProcess(),
111	n/a	handle, proc, access, False, 0)
112	n/a	finally:
113	n/a	_winapi.CloseHandle(proc)
114	n/a	self._access = access
115	n/a	self._pid = pid
116	n/a
117	n/a	def detach(self):
118	n/a	'''Get the handle. This should only be called once.'''
119	n/a	# retrieve handle from process which currently owns it
120	n/a	if self._pid == os.getpid():
121	n/a	# The handle has already been duplicated for this process.
122	n/a	return self._handle
123	n/a	# We must steal the handle from the process whose pid is self._pid.
124	n/a	proc = _winapi.OpenProcess(_winapi.PROCESS_DUP_HANDLE, False,
125	n/a	self._pid)
126	n/a	try:
127	n/a	return _winapi.DuplicateHandle(
128	n/a	proc, self._handle, _winapi.GetCurrentProcess(),
129	n/a	self._access, False, _winapi.DUPLICATE_CLOSE_SOURCE)
130	n/a	finally:
131	n/a	_winapi.CloseHandle(proc)
132	n/a
133	n/a	else:
134	n/a	# Unix
135	n/a	__all__ += ['DupFd', 'sendfds', 'recvfds']
136	n/a	import array
137	n/a
138	n/a	# On MacOSX we should acknowledge receipt of fds -- see Issue14669
139	n/a	ACKNOWLEDGE = sys.platform == 'darwin'
140	n/a
141	n/a	def sendfds(sock, fds):
142	n/a	'''Send an array of fds over an AF_UNIX socket.'''
143	n/a	fds = array.array('i', fds)
144	n/a	msg = bytes([len(fds) % 256])
145	n/a	sock.sendmsg([msg], [(socket.SOL_SOCKET, socket.SCM_RIGHTS, fds)])
146	n/a	if ACKNOWLEDGE and sock.recv(1) != b'A':
147	n/a	raise RuntimeError('did not receive acknowledgement of fd')
148	n/a
149	n/a	def recvfds(sock, size):
150	n/a	'''Receive an array of fds over an AF_UNIX socket.'''
151	n/a	a = array.array('i')
152	n/a	bytes_size = a.itemsize * size
153	n/a	msg, ancdata, flags, addr = sock.recvmsg(1, socket.CMSG_LEN(bytes_size))
154	n/a	if not msg and not ancdata:
155	n/a	raise EOFError
156	n/a	try:
157	n/a	if ACKNOWLEDGE:
158	n/a	sock.send(b'A')
159	n/a	if len(ancdata) != 1:
160	n/a	raise RuntimeError('received %d items of ancdata' %
161	n/a	len(ancdata))
162	n/a	cmsg_level, cmsg_type, cmsg_data = ancdata[0]
163	n/a	if (cmsg_level == socket.SOL_SOCKET and
164	n/a	cmsg_type == socket.SCM_RIGHTS):
165	n/a	if len(cmsg_data) % a.itemsize != 0:
166	n/a	raise ValueError
167	n/a	a.frombytes(cmsg_data)
168	n/a	assert len(a) % 256 == msg[0]
169	n/a	return list(a)
170	n/a	except (ValueError, IndexError):
171	n/a	pass
172	n/a	raise RuntimeError('Invalid data received')
173	n/a
174	n/a	def send_handle(conn, handle, destination_pid):
175	n/a	'''Send a handle over a local connection.'''
176	n/a	with socket.fromfd(conn.fileno(), socket.AF_UNIX, socket.SOCK_STREAM) as s:
177	n/a	sendfds(s, [handle])
178	n/a
179	n/a	def recv_handle(conn):
180	n/a	'''Receive a handle over a local connection.'''
181	n/a	with socket.fromfd(conn.fileno(), socket.AF_UNIX, socket.SOCK_STREAM) as s:
182	n/a	return recvfds(s, 1)[0]
183	n/a
184	n/a	def DupFd(fd):
185	n/a	'''Return a wrapper for an fd.'''
186	n/a	popen_obj = context.get_spawning_popen()
187	n/a	if popen_obj is not None:
188	n/a	return popen_obj.DupFd(popen_obj.duplicate_for_child(fd))
189	n/a	elif HAVE_SEND_HANDLE:
190	n/a	from . import resource_sharer
191	n/a	return resource_sharer.DupFd(fd)
192	n/a	else:
193	n/a	raise ValueError('SCM_RIGHTS appears not to be available')
194	n/a
195	n/a	#
196	n/a	# Try making some callable types picklable
197	n/a	#
198	n/a
199	n/a	def _reduce_method(m):
200	n/a	if m.__self__ is None:
201	n/a	return getattr, (m.__class__, m.__func__.__name__)
202	n/a	else:
203	n/a	return getattr, (m.__self__, m.__func__.__name__)
204	n/a	class _C:
205	n/a	def f(self):
206	n/a	pass
207	n/a	register(type(_C().f), _reduce_method)
208	n/a
209	n/a
210	n/a	def _reduce_method_descriptor(m):
211	n/a	return getattr, (m.__objclass__, m.__name__)
212	n/a	register(type(list.append), _reduce_method_descriptor)
213	n/a	register(type(int.__add__), _reduce_method_descriptor)
214	n/a
215	n/a
216	n/a	def _reduce_partial(p):
217	n/a	return _rebuild_partial, (p.func, p.args, p.keywords or {})
218	n/a	def _rebuild_partial(func, args, keywords):
219	n/a	return functools.partial(func, args, *keywords)
220	n/a	register(functools.partial, _reduce_partial)
221	n/a
222	n/a	#
223	n/a	# Make sockets picklable
224	n/a	#
225	n/a
226	n/a	if sys.platform == 'win32':
227	n/a	def _reduce_socket(s):
228	n/a	from .resource_sharer import DupSocket
229	n/a	return _rebuild_socket, (DupSocket(s),)
230	n/a	def _rebuild_socket(ds):
231	n/a	return ds.detach()
232	n/a	register(socket.socket, _reduce_socket)
233	n/a
234	n/a	else:
235	n/a	def _reduce_socket(s):
236	n/a	df = DupFd(s.fileno())
237	n/a	return _rebuild_socket, (df, s.family, s.type, s.proto)
238	n/a	def _rebuild_socket(df, family, type, proto):
239	n/a	fd = df.detach()
240	n/a	return socket.socket(family, type, proto, fileno=fd)
241	n/a	register(socket.socket, _reduce_socket)
242	n/a
243	n/a
244	n/a	class AbstractReducer(metaclass=ABCMeta):
245	n/a	'''Abstract base class for use in implementing a Reduction class
246	n/a	suitable for use in replacing the standard reduction mechanism
247	n/a	used in multiprocessing.'''
248	n/a	ForkingPickler = ForkingPickler
249	n/a	register = register
250	n/a	dump = dump
251	n/a	send_handle = send_handle
252	n/a	recv_handle = recv_handle
253	n/a
254	n/a	if sys.platform == 'win32':
255	n/a	steal_handle = steal_handle
256	n/a	duplicate = duplicate
257	n/a	DupHandle = DupHandle
258	n/a	else:
259	n/a	sendfds = sendfds
260	n/a	recvfds = recvfds
261	n/a	DupFd = DupFd
262	n/a
263	n/a	_reduce_method = _reduce_method
264	n/a	_reduce_method_descriptor = _reduce_method_descriptor
265	n/a	_rebuild_partial = _rebuild_partial
266	n/a	_reduce_socket = _reduce_socket
267	n/a	_rebuild_socket = _rebuild_socket
268	n/a
269	n/a	def __init__(self, *args):
270	n/a	register(type(_C().f), _reduce_method)
271	n/a	register(type(list.append), _reduce_method_descriptor)
272	n/a	register(type(int.__add__), _reduce_method_descriptor)
273	n/a	register(functools.partial, _reduce_partial)
274	n/a	register(socket.socket, _reduce_socket)