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