Mercurial > repos > rliterman > csp2

annotate CSP2/CSP2_env/env-d9b9114564458d9d-741b3de822f2aaca6c6caa4325c4afce/lib/python3.8/multiprocessing/reduction.py @ 68:5028fdace37b

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