File size: 6,266 Bytes
d1ceb73 |
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 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 |
"""AsyncIO support for zmq
Requires asyncio and Python 3.
"""
# Copyright (c) PyZMQ Developers.
# Distributed under the terms of the Modified BSD License.
from __future__ import annotations
import asyncio
import selectors
import sys
import warnings
from asyncio import Future, SelectorEventLoop
from weakref import WeakKeyDictionary
import zmq as _zmq
from zmq import _future
# registry of asyncio loop : selector thread
_selectors: WeakKeyDictionary = WeakKeyDictionary()
class ProactorSelectorThreadWarning(RuntimeWarning):
"""Warning class for notifying about the extra thread spawned by tornado
We automatically support proactor via tornado's AddThreadSelectorEventLoop"""
def _get_selector_windows(
asyncio_loop,
) -> asyncio.AbstractEventLoop:
"""Get selector-compatible loop
Returns an object with ``add_reader`` family of methods,
either the loop itself or a SelectorThread instance.
Workaround Windows proactor removal of
*reader methods, which we need for zmq sockets.
"""
if asyncio_loop in _selectors:
return _selectors[asyncio_loop]
# detect add_reader instead of checking for proactor?
if hasattr(asyncio, "ProactorEventLoop") and isinstance(
asyncio_loop,
asyncio.ProactorEventLoop, # type: ignore
):
try:
from tornado.platform.asyncio import AddThreadSelectorEventLoop
except ImportError:
raise RuntimeError(
"Proactor event loop does not implement add_reader family of methods required for zmq."
" zmq will work with proactor if tornado >= 6.1 can be found."
" Use `asyncio.set_event_loop_policy(WindowsSelectorEventLoopPolicy())`"
" or install 'tornado>=6.1' to avoid this error."
)
warnings.warn(
"Proactor event loop does not implement add_reader family of methods required for zmq."
" Registering an additional selector thread for add_reader support via tornado."
" Use `asyncio.set_event_loop_policy(WindowsSelectorEventLoopPolicy())`"
" to avoid this warning.",
RuntimeWarning,
# stacklevel 5 matches most likely zmq.asyncio.Context().socket()
stacklevel=5,
)
selector_loop = _selectors[asyncio_loop] = AddThreadSelectorEventLoop(
asyncio_loop
) # type: ignore
# patch loop.close to also close the selector thread
loop_close = asyncio_loop.close
def _close_selector_and_loop():
# restore original before calling selector.close,
# which in turn calls eventloop.close!
asyncio_loop.close = loop_close
_selectors.pop(asyncio_loop, None)
selector_loop.close()
asyncio_loop.close = _close_selector_and_loop # type: ignore # mypy bug - assign a function to method
return selector_loop
else:
return asyncio_loop
def _get_selector_noop(loop) -> asyncio.AbstractEventLoop:
"""no-op on non-Windows"""
return loop
if sys.platform == "win32":
_get_selector = _get_selector_windows
else:
_get_selector = _get_selector_noop
class _AsyncIO:
_Future = Future
_WRITE = selectors.EVENT_WRITE
_READ = selectors.EVENT_READ
def _default_loop(self):
try:
return asyncio.get_running_loop()
except RuntimeError:
warnings.warn(
"No running event loop. zmq.asyncio should be used from within an asyncio loop.",
RuntimeWarning,
stacklevel=4,
)
# get_event_loop deprecated in 3.10:
return asyncio.get_event_loop()
class Poller(_AsyncIO, _future._AsyncPoller):
"""Poller returning asyncio.Future for poll results."""
def _watch_raw_socket(self, loop, socket, evt, f):
"""Schedule callback for a raw socket"""
selector = _get_selector(loop)
if evt & self._READ:
selector.add_reader(socket, lambda *args: f())
if evt & self._WRITE:
selector.add_writer(socket, lambda *args: f())
def _unwatch_raw_sockets(self, loop, *sockets):
"""Unschedule callback for a raw socket"""
selector = _get_selector(loop)
for socket in sockets:
selector.remove_reader(socket)
selector.remove_writer(socket)
class Socket(_AsyncIO, _future._AsyncSocket):
"""Socket returning asyncio Futures for send/recv/poll methods."""
_poller_class = Poller
def _get_selector(self, io_loop=None):
if io_loop is None:
io_loop = self._get_loop()
return _get_selector(io_loop)
def _init_io_state(self, io_loop=None):
"""initialize the ioloop event handler"""
self._get_selector(io_loop).add_reader(
self._fd, lambda: self._handle_events(0, 0)
)
def _clear_io_state(self):
"""clear any ioloop event handler
called once at close
"""
loop = self._current_loop
if loop and not loop.is_closed() and self._fd != -1:
self._get_selector(loop).remove_reader(self._fd)
Poller._socket_class = Socket
class Context(_zmq.Context[Socket]):
"""Context for creating asyncio-compatible Sockets"""
_socket_class = Socket
# avoid sharing instance with base Context class
_instance = None
class ZMQEventLoop(SelectorEventLoop):
"""DEPRECATED: AsyncIO eventloop using zmq_poll.
pyzmq sockets should work with any asyncio event loop as of pyzmq 17.
"""
def __init__(self, selector=None):
_deprecated()
return super().__init__(selector)
_loop = None
def _deprecated():
if _deprecated.called: # type: ignore
return
_deprecated.called = True # type: ignore
warnings.warn(
"ZMQEventLoop and zmq.asyncio.install are deprecated in pyzmq 17. Special eventloop integration is no longer needed.",
DeprecationWarning,
stacklevel=3,
)
_deprecated.called = False # type: ignore
def install():
"""DEPRECATED: No longer needed in pyzmq 17"""
_deprecated()
__all__ = [
"Context",
"Socket",
"Poller",
"ZMQEventLoop",
"install",
]
|