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