1###############################################################################
2# Extra reducers for Unix based system and connections objects
3#
4# author: Thomas Moreau and Olivier Grisel
5#
6# adapted from multiprocessing/reduction.py (17/02/2017)
7# * Add adapted reduction for LokyProcesses and socket/Connection
8#
9import os
10import socket
11import _socket
12from multiprocessing.connection import Connection
13from multiprocessing.context import get_spawning_popen
14
15from .reduction import register
16
17HAVE_SEND_HANDLE = (
18 hasattr(socket, "CMSG_LEN")
19 and hasattr(socket, "SCM_RIGHTS")
20 and hasattr(socket.socket, "sendmsg")
21)
22
23
24def _mk_inheritable(fd):
25 os.set_inheritable(fd, True)
26 return fd
27
28
29def DupFd(fd):
30 """Return a wrapper for an fd."""
31 popen_obj = get_spawning_popen()
32 if popen_obj is not None:
33 return popen_obj.DupFd(popen_obj.duplicate_for_child(fd))
34 elif HAVE_SEND_HANDLE:
35 from multiprocessing import resource_sharer
36
37 return resource_sharer.DupFd(fd)
38 else:
39 raise TypeError(
40 "Cannot pickle connection object. This object can only be "
41 "passed when spawning a new process"
42 )
43
44
45def _reduce_socket(s):
46 df = DupFd(s.fileno())
47 return _rebuild_socket, (df, s.family, s.type, s.proto)
48
49
50def _rebuild_socket(df, family, type, proto):
51 fd = df.detach()
52 return socket.fromfd(fd, family, type, proto)
53
54
55def rebuild_connection(df, readable, writable):
56 fd = df.detach()
57 return Connection(fd, readable, writable)
58
59
60def reduce_connection(conn):
61 df = DupFd(conn.fileno())
62 return rebuild_connection, (df, conn.readable, conn.writable)
63
64
65register(socket.socket, _reduce_socket)
66register(_socket.socket, _reduce_socket)
67register(Connection, reduce_connection)