windows_utils.py 4.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173
  1. """Various Windows specific bits and pieces."""
  2. import sys
  3. if sys.platform != 'win32': # pragma: no cover
  4. raise ImportError('win32 only')
  5. import _winapi
  6. import itertools
  7. import msvcrt
  8. import os
  9. import subprocess
  10. import tempfile
  11. import warnings
  12. __all__ = 'pipe', 'Popen', 'PIPE', 'PipeHandle'
  13. # Constants/globals
  14. BUFSIZE = 8192
  15. PIPE = subprocess.PIPE
  16. STDOUT = subprocess.STDOUT
  17. _mmap_counter = itertools.count()
  18. # Replacement for os.pipe() using handles instead of fds
  19. def pipe(*, duplex=False, overlapped=(True, True), bufsize=BUFSIZE):
  20. """Like os.pipe() but with overlapped support and using handles not fds."""
  21. address = tempfile.mktemp(
  22. prefix=r'\\.\pipe\python-pipe-{:d}-{:d}-'.format(
  23. os.getpid(), next(_mmap_counter)))
  24. if duplex:
  25. openmode = _winapi.PIPE_ACCESS_DUPLEX
  26. access = _winapi.GENERIC_READ | _winapi.GENERIC_WRITE
  27. obsize, ibsize = bufsize, bufsize
  28. else:
  29. openmode = _winapi.PIPE_ACCESS_INBOUND
  30. access = _winapi.GENERIC_WRITE
  31. obsize, ibsize = 0, bufsize
  32. openmode |= _winapi.FILE_FLAG_FIRST_PIPE_INSTANCE
  33. if overlapped[0]:
  34. openmode |= _winapi.FILE_FLAG_OVERLAPPED
  35. if overlapped[1]:
  36. flags_and_attribs = _winapi.FILE_FLAG_OVERLAPPED
  37. else:
  38. flags_and_attribs = 0
  39. h1 = h2 = None
  40. try:
  41. h1 = _winapi.CreateNamedPipe(
  42. address, openmode, _winapi.PIPE_WAIT,
  43. 1, obsize, ibsize, _winapi.NMPWAIT_WAIT_FOREVER, _winapi.NULL)
  44. h2 = _winapi.CreateFile(
  45. address, access, 0, _winapi.NULL, _winapi.OPEN_EXISTING,
  46. flags_and_attribs, _winapi.NULL)
  47. ov = _winapi.ConnectNamedPipe(h1, overlapped=True)
  48. ov.GetOverlappedResult(True)
  49. return h1, h2
  50. except:
  51. if h1 is not None:
  52. _winapi.CloseHandle(h1)
  53. if h2 is not None:
  54. _winapi.CloseHandle(h2)
  55. raise
  56. # Wrapper for a pipe handle
  57. class PipeHandle:
  58. """Wrapper for an overlapped pipe handle which is vaguely file-object like.
  59. The IOCP event loop can use these instead of socket objects.
  60. """
  61. def __init__(self, handle):
  62. self._handle = handle
  63. def __repr__(self):
  64. if self._handle is not None:
  65. handle = f'handle={self._handle!r}'
  66. else:
  67. handle = 'closed'
  68. return f'<{self.__class__.__name__} {handle}>'
  69. @property
  70. def handle(self):
  71. return self._handle
  72. def fileno(self):
  73. if self._handle is None:
  74. raise ValueError("I/O operation on closed pipe")
  75. return self._handle
  76. def close(self, *, CloseHandle=_winapi.CloseHandle):
  77. if self._handle is not None:
  78. CloseHandle(self._handle)
  79. self._handle = None
  80. def __del__(self, _warn=warnings.warn):
  81. if self._handle is not None:
  82. _warn(f"unclosed {self!r}", ResourceWarning, source=self)
  83. self.close()
  84. def __enter__(self):
  85. return self
  86. def __exit__(self, t, v, tb):
  87. self.close()
  88. # Replacement for subprocess.Popen using overlapped pipe handles
  89. class Popen(subprocess.Popen):
  90. """Replacement for subprocess.Popen using overlapped pipe handles.
  91. The stdin, stdout, stderr are None or instances of PipeHandle.
  92. """
  93. def __init__(self, args, stdin=None, stdout=None, stderr=None, **kwds):
  94. assert not kwds.get('universal_newlines')
  95. assert kwds.get('bufsize', 0) == 0
  96. stdin_rfd = stdout_wfd = stderr_wfd = None
  97. stdin_wh = stdout_rh = stderr_rh = None
  98. if stdin == PIPE:
  99. stdin_rh, stdin_wh = pipe(overlapped=(False, True), duplex=True)
  100. stdin_rfd = msvcrt.open_osfhandle(stdin_rh, os.O_RDONLY)
  101. else:
  102. stdin_rfd = stdin
  103. if stdout == PIPE:
  104. stdout_rh, stdout_wh = pipe(overlapped=(True, False))
  105. stdout_wfd = msvcrt.open_osfhandle(stdout_wh, 0)
  106. else:
  107. stdout_wfd = stdout
  108. if stderr == PIPE:
  109. stderr_rh, stderr_wh = pipe(overlapped=(True, False))
  110. stderr_wfd = msvcrt.open_osfhandle(stderr_wh, 0)
  111. elif stderr == STDOUT:
  112. stderr_wfd = stdout_wfd
  113. else:
  114. stderr_wfd = stderr
  115. try:
  116. super().__init__(args, stdin=stdin_rfd, stdout=stdout_wfd,
  117. stderr=stderr_wfd, **kwds)
  118. except:
  119. for h in (stdin_wh, stdout_rh, stderr_rh):
  120. if h is not None:
  121. _winapi.CloseHandle(h)
  122. raise
  123. else:
  124. if stdin_wh is not None:
  125. self.stdin = PipeHandle(stdin_wh)
  126. if stdout_rh is not None:
  127. self.stdout = PipeHandle(stdout_rh)
  128. if stderr_rh is not None:
  129. self.stderr = PipeHandle(stderr_rh)
  130. finally:
  131. if stdin == PIPE:
  132. os.close(stdin_rfd)
  133. if stdout == PIPE:
  134. os.close(stdout_wfd)
  135. if stderr == PIPE:
  136. os.close(stderr_wfd)