win32.py 5.0 KB

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