Development of an internal social media platform with personalised dashboards for students
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

reduction.py 7.8KB

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