summaryrefslogtreecommitdiffstats
path: root/src/console/zmq/utils/win32.py
blob: ea75829912ba144886c4d5fc86e0f10a2ffc7004 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
"""Win32 compatibility utilities."""

#-----------------------------------------------------------------------------
# Copyright (C) PyZMQ Developers
# Distributed under the terms of the Modified BSD License.
#-----------------------------------------------------------------------------

import os

# No-op implementation for other platforms.
class _allow_interrupt(object):
    """Utility for fixing CTRL-C events on Windows.

    On Windows, the Python interpreter intercepts CTRL-C events in order to
    translate them into ``KeyboardInterrupt`` exceptions.  It (presumably)
    does this by setting a flag in its "control control handler" and
    checking it later at a convenient location in the interpreter.

    However, when the Python interpreter is blocked waiting for the ZMQ
    poll operation to complete, it must wait for ZMQ's ``select()``
    operation to complete before translating the CTRL-C event into the
    ``KeyboardInterrupt`` exception.

    The only way to fix this seems to be to add our own "console control
    handler" and perform some application-defined operation that will
    unblock the ZMQ polling operation in order to force ZMQ to pass control
    back to the Python interpreter.

    This context manager performs all that Windows-y stuff, providing you
    with a hook that is called when a CTRL-C event is intercepted.  This
    hook allows you to unblock your ZMQ poll operation immediately, which
    will then result in the expected ``KeyboardInterrupt`` exception.

    Without this context manager, your ZMQ-based application will not
    respond normally to CTRL-C events on Windows.  If a CTRL-C event occurs
    while blocked on ZMQ socket polling, the translation to a
    ``KeyboardInterrupt`` exception will be delayed until the I/O completes
    and control returns to the Python interpreter (this may never happen if
    you use an infinite timeout).

    A no-op implementation is provided on non-Win32 systems to avoid the
    application from having to conditionally use it.

    Example usage:

    .. sourcecode:: python

       def stop_my_application():
           # ...

       with allow_interrupt(stop_my_application):
           # main polling loop.

    In a typical ZMQ application, you would use the "self pipe trick" to
    send message to a ``PAIR`` socket in order to interrupt your blocking
    socket polling operation.

    In a Tornado event loop, you can use the ``IOLoop.stop`` method to
    unblock your I/O loop.
    """

    def __init__(self, action=None):
        """Translate ``action`` into a CTRL-C handler.

        ``action`` is a callable that takes no arguments and returns no
        value (returned value is ignored).  It must *NEVER* raise an
        exception.
        
        If unspecified, a no-op will be used.
        """
        self._init_action(action)
    
    def _init_action(self, action):
        pass

    def __enter__(self):
        return self

    def __exit__(self, *args):
        return

if os.name == 'nt':
    from ctypes import WINFUNCTYPE, windll
    from ctypes.wintypes import BOOL, DWORD

    kernel32 = windll.LoadLibrary('kernel32')

    # <http://msdn.microsoft.com/en-us/library/ms686016.aspx>
    PHANDLER_ROUTINE = WINFUNCTYPE(BOOL, DWORD)
    SetConsoleCtrlHandler = kernel32.SetConsoleCtrlHandler
    SetConsoleCtrlHandler.argtypes = (PHANDLER_ROUTINE, BOOL)
    SetConsoleCtrlHandler.restype = BOOL

    class allow_interrupt(_allow_interrupt):
        __doc__ = _allow_interrupt.__doc__

        def _init_action(self, action):
            if action is None:
                action = lambda: None
            self.action = action
            @PHANDLER_ROUTINE
            def handle(event):
                if event == 0:  # CTRL_C_EVENT
                    action()
                    # Typical C implementations would return 1 to indicate that
                    # the event was processed and other control handlers in the
                    # stack should not be executed.  However, that would
                    # prevent the Python interpreter's handler from translating
                    # CTRL-C to a `KeyboardInterrupt` exception, so we pretend
                    # that we didn't handle it.
                return 0
            self.handle = handle

        def __enter__(self):
            """Install the custom CTRL-C handler."""
            result = SetConsoleCtrlHandler(self.handle, 1)
            if result == 0:
                # Have standard library automatically call `GetLastError()` and
                # `FormatMessage()` into a nice exception object :-)
                raise WindowsError()

        def __exit__(self, *args):
            """Remove the custom CTRL-C handler."""
            result = SetConsoleCtrlHandler(self.handle, 0)
            if result == 0:
                # Have standard library automatically call `GetLastError()` and
                # `FormatMessage()` into a nice exception object :-)
                raise WindowsError()
else:
    class allow_interrupt(_allow_interrupt):
        __doc__ = _allow_interrupt.__doc__
        pass