win32.py 4.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130
  1. """Win32 compatibility utilities."""
  2. # -----------------------------------------------------------------------------
  3. # Copyright (C) PyZMQ Developers
  4. # Distributed under the terms of the Modified BSD License.
  5. # -----------------------------------------------------------------------------
  6. from __future__ import annotations
  7. import os
  8. from typing import Any, Callable
  9. class allow_interrupt:
  10. """Utility for fixing CTRL-C events on Windows.
  11. On Windows, the Python interpreter intercepts CTRL-C events in order to
  12. translate them into ``KeyboardInterrupt`` exceptions. It (presumably)
  13. does this by setting a flag in its "console control handler" and
  14. checking it later at a convenient location in the interpreter.
  15. However, when the Python interpreter is blocked waiting for the ZMQ
  16. poll operation to complete, it must wait for ZMQ's ``select()``
  17. operation to complete before translating the CTRL-C event into the
  18. ``KeyboardInterrupt`` exception.
  19. The only way to fix this seems to be to add our own "console control
  20. handler" and perform some application-defined operation that will
  21. unblock the ZMQ polling operation in order to force ZMQ to pass control
  22. back to the Python interpreter.
  23. This context manager performs all that Windows-y stuff, providing you
  24. with a hook that is called when a CTRL-C event is intercepted. This
  25. hook allows you to unblock your ZMQ poll operation immediately, which
  26. will then result in the expected ``KeyboardInterrupt`` exception.
  27. Without this context manager, your ZMQ-based application will not
  28. respond normally to CTRL-C events on Windows. If a CTRL-C event occurs
  29. while blocked on ZMQ socket polling, the translation to a
  30. ``KeyboardInterrupt`` exception will be delayed until the I/O completes
  31. and control returns to the Python interpreter (this may never happen if
  32. you use an infinite timeout).
  33. A no-op implementation is provided on non-Win32 systems to avoid the
  34. application from having to conditionally use it.
  35. Example usage:
  36. .. sourcecode:: python
  37. def stop_my_application():
  38. # ...
  39. with allow_interrupt(stop_my_application):
  40. # main polling loop.
  41. In a typical ZMQ application, you would use the "self pipe trick" to
  42. send message to a ``PAIR`` socket in order to interrupt your blocking
  43. socket polling operation.
  44. In a Tornado event loop, you can use the ``IOLoop.stop`` method to
  45. unblock your I/O loop.
  46. """
  47. def __init__(self, action: Callable[[], Any] | None = None) -> None:
  48. """Translate ``action`` into a CTRL-C handler.
  49. ``action`` is a callable that takes no arguments and returns no
  50. value (returned value is ignored). It must *NEVER* raise an
  51. exception.
  52. If unspecified, a no-op will be used.
  53. """
  54. if os.name != "nt":
  55. return
  56. self._init_action(action)
  57. def _init_action(self, action):
  58. from ctypes import WINFUNCTYPE, windll
  59. from ctypes.wintypes import BOOL, DWORD
  60. kernel32 = windll.LoadLibrary('kernel32')
  61. # <http://msdn.microsoft.com/en-us/library/ms686016.aspx>
  62. PHANDLER_ROUTINE = WINFUNCTYPE(BOOL, DWORD)
  63. SetConsoleCtrlHandler = self._SetConsoleCtrlHandler = (
  64. kernel32.SetConsoleCtrlHandler
  65. )
  66. SetConsoleCtrlHandler.argtypes = (PHANDLER_ROUTINE, BOOL)
  67. SetConsoleCtrlHandler.restype = BOOL
  68. if action is None:
  69. def action():
  70. return None
  71. self.action = action
  72. @PHANDLER_ROUTINE
  73. def handle(event):
  74. if event == 0: # CTRL_C_EVENT
  75. action()
  76. # Typical C implementations would return 1 to indicate that
  77. # the event was processed and other control handlers in the
  78. # stack should not be executed. However, that would
  79. # prevent the Python interpreter's handler from translating
  80. # CTRL-C to a `KeyboardInterrupt` exception, so we pretend
  81. # that we didn't handle it.
  82. return 0
  83. self.handle = handle
  84. def __enter__(self):
  85. """Install the custom CTRL-C handler."""
  86. if os.name != "nt":
  87. return
  88. result = self._SetConsoleCtrlHandler(self.handle, 1)
  89. if result == 0:
  90. # Have standard library automatically call `GetLastError()` and
  91. # `FormatMessage()` into a nice exception object :-)
  92. raise OSError()
  93. def __exit__(self, *args):
  94. """Remove the custom CTRL-C handler."""
  95. if os.name != "nt":
  96. return
  97. result = self._SetConsoleCtrlHandler(self.handle, 0)
  98. if result == 0:
  99. # Have standard library automatically call `GetLastError()` and
  100. # `FormatMessage()` into a nice exception object :-)
  101. raise OSError()