windows_utils.py 6.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223
  1. """
  2. Various Windows specific bits and pieces
  3. """
  4. import sys
  5. if sys.platform != 'win32': # pragma: no cover
  6. raise ImportError('win32 only')
  7. import _winapi
  8. import itertools
  9. import msvcrt
  10. import os
  11. import socket
  12. import subprocess
  13. import tempfile
  14. import warnings
  15. __all__ = ['socketpair', 'pipe', 'Popen', 'PIPE', 'PipeHandle']
  16. # Constants/globals
  17. BUFSIZE = 8192
  18. PIPE = subprocess.PIPE
  19. STDOUT = subprocess.STDOUT
  20. _mmap_counter = itertools.count()
  21. if hasattr(socket, 'socketpair'):
  22. # Since Python 3.5, socket.socketpair() is now also available on Windows
  23. socketpair = socket.socketpair
  24. else:
  25. # Replacement for socket.socketpair()
  26. def socketpair(family=socket.AF_INET, type=socket.SOCK_STREAM, proto=0):
  27. """A socket pair usable as a self-pipe, for Windows.
  28. Origin: https://gist.github.com/4325783, by Geert Jansen.
  29. Public domain.
  30. """
  31. if family == socket.AF_INET:
  32. host = '127.0.0.1'
  33. elif family == socket.AF_INET6:
  34. host = '::1'
  35. else:
  36. raise ValueError("Only AF_INET and AF_INET6 socket address "
  37. "families are supported")
  38. if type != socket.SOCK_STREAM:
  39. raise ValueError("Only SOCK_STREAM socket type is supported")
  40. if proto != 0:
  41. raise ValueError("Only protocol zero is supported")
  42. # We create a connected TCP socket. Note the trick with setblocking(0)
  43. # that prevents us from having to create a thread.
  44. lsock = socket.socket(family, type, proto)
  45. try:
  46. lsock.bind((host, 0))
  47. lsock.listen(1)
  48. # On IPv6, ignore flow_info and scope_id
  49. addr, port = lsock.getsockname()[:2]
  50. csock = socket.socket(family, type, proto)
  51. try:
  52. csock.setblocking(False)
  53. try:
  54. csock.connect((addr, port))
  55. except (BlockingIOError, InterruptedError):
  56. pass
  57. csock.setblocking(True)
  58. ssock, _ = lsock.accept()
  59. except:
  60. csock.close()
  61. raise
  62. finally:
  63. lsock.close()
  64. return (ssock, csock)
  65. # Replacement for os.pipe() using handles instead of fds
  66. def pipe(*, duplex=False, overlapped=(True, True), bufsize=BUFSIZE):
  67. """Like os.pipe() but with overlapped support and using handles not fds."""
  68. address = tempfile.mktemp(prefix=r'\\.\pipe\python-pipe-%d-%d-' %
  69. (os.getpid(), next(_mmap_counter)))
  70. if duplex:
  71. openmode = _winapi.PIPE_ACCESS_DUPLEX
  72. access = _winapi.GENERIC_READ | _winapi.GENERIC_WRITE
  73. obsize, ibsize = bufsize, bufsize
  74. else:
  75. openmode = _winapi.PIPE_ACCESS_INBOUND
  76. access = _winapi.GENERIC_WRITE
  77. obsize, ibsize = 0, bufsize
  78. openmode |= _winapi.FILE_FLAG_FIRST_PIPE_INSTANCE
  79. if overlapped[0]:
  80. openmode |= _winapi.FILE_FLAG_OVERLAPPED
  81. if overlapped[1]:
  82. flags_and_attribs = _winapi.FILE_FLAG_OVERLAPPED
  83. else:
  84. flags_and_attribs = 0
  85. h1 = h2 = None
  86. try:
  87. h1 = _winapi.CreateNamedPipe(
  88. address, openmode, _winapi.PIPE_WAIT,
  89. 1, obsize, ibsize, _winapi.NMPWAIT_WAIT_FOREVER, _winapi.NULL)
  90. h2 = _winapi.CreateFile(
  91. address, access, 0, _winapi.NULL, _winapi.OPEN_EXISTING,
  92. flags_and_attribs, _winapi.NULL)
  93. ov = _winapi.ConnectNamedPipe(h1, overlapped=True)
  94. ov.GetOverlappedResult(True)
  95. return h1, h2
  96. except:
  97. if h1 is not None:
  98. _winapi.CloseHandle(h1)
  99. if h2 is not None:
  100. _winapi.CloseHandle(h2)
  101. raise
  102. # Wrapper for a pipe handle
  103. class PipeHandle:
  104. """Wrapper for an overlapped pipe handle which is vaguely file-object like.
  105. The IOCP event loop can use these instead of socket objects.
  106. """
  107. def __init__(self, handle):
  108. self._handle = handle
  109. def __repr__(self):
  110. if self._handle is not None:
  111. handle = 'handle=%r' % self._handle
  112. else:
  113. handle = 'closed'
  114. return '<%s %s>' % (self.__class__.__name__, handle)
  115. @property
  116. def handle(self):
  117. return self._handle
  118. def fileno(self):
  119. if self._handle is None:
  120. raise ValueError("I/O operatioon on closed pipe")
  121. return self._handle
  122. def close(self, *, CloseHandle=_winapi.CloseHandle):
  123. if self._handle is not None:
  124. CloseHandle(self._handle)
  125. self._handle = None
  126. def __del__(self):
  127. if self._handle is not None:
  128. warnings.warn("unclosed %r" % self, ResourceWarning)
  129. self.close()
  130. def __enter__(self):
  131. return self
  132. def __exit__(self, t, v, tb):
  133. self.close()
  134. # Replacement for subprocess.Popen using overlapped pipe handles
  135. class Popen(subprocess.Popen):
  136. """Replacement for subprocess.Popen using overlapped pipe handles.
  137. The stdin, stdout, stderr are None or instances of PipeHandle.
  138. """
  139. def __init__(self, args, stdin=None, stdout=None, stderr=None, **kwds):
  140. assert not kwds.get('universal_newlines')
  141. assert kwds.get('bufsize', 0) == 0
  142. stdin_rfd = stdout_wfd = stderr_wfd = None
  143. stdin_wh = stdout_rh = stderr_rh = None
  144. if stdin == PIPE:
  145. stdin_rh, stdin_wh = pipe(overlapped=(False, True), duplex=True)
  146. stdin_rfd = msvcrt.open_osfhandle(stdin_rh, os.O_RDONLY)
  147. else:
  148. stdin_rfd = stdin
  149. if stdout == PIPE:
  150. stdout_rh, stdout_wh = pipe(overlapped=(True, False))
  151. stdout_wfd = msvcrt.open_osfhandle(stdout_wh, 0)
  152. else:
  153. stdout_wfd = stdout
  154. if stderr == PIPE:
  155. stderr_rh, stderr_wh = pipe(overlapped=(True, False))
  156. stderr_wfd = msvcrt.open_osfhandle(stderr_wh, 0)
  157. elif stderr == STDOUT:
  158. stderr_wfd = stdout_wfd
  159. else:
  160. stderr_wfd = stderr
  161. try:
  162. super().__init__(args, stdin=stdin_rfd, stdout=stdout_wfd,
  163. stderr=stderr_wfd, **kwds)
  164. except:
  165. for h in (stdin_wh, stdout_rh, stderr_rh):
  166. if h is not None:
  167. _winapi.CloseHandle(h)
  168. raise
  169. else:
  170. if stdin_wh is not None:
  171. self.stdin = PipeHandle(stdin_wh)
  172. if stdout_rh is not None:
  173. self.stdout = PipeHandle(stdout_rh)
  174. if stderr_rh is not None:
  175. self.stderr = PipeHandle(stderr_rh)
  176. finally:
  177. if stdin == PIPE:
  178. os.close(stdin_rfd)
  179. if stdout == PIPE:
  180. os.close(stdout_wfd)
  181. if stderr == PIPE:
  182. os.close(stderr_wfd)