|
"""0MQ Socket pure Python methods.""" |
|
|
|
|
|
|
|
|
|
from __future__ import annotations |
|
|
|
import errno |
|
import pickle |
|
import random |
|
import sys |
|
from typing import ( |
|
Any, |
|
Callable, |
|
Generic, |
|
List, |
|
Sequence, |
|
TypeVar, |
|
Union, |
|
cast, |
|
overload, |
|
) |
|
from warnings import warn |
|
|
|
import zmq |
|
from zmq._typing import Literal, TypeAlias |
|
from zmq.backend import Socket as SocketBase |
|
from zmq.error import ZMQBindError, ZMQError |
|
from zmq.utils import jsonapi |
|
from zmq.utils.interop import cast_int_addr |
|
|
|
from ..constants import SocketOption, SocketType, _OptType |
|
from .attrsettr import AttributeSetter |
|
from .poll import Poller |
|
|
|
try: |
|
DEFAULT_PROTOCOL = pickle.DEFAULT_PROTOCOL |
|
except AttributeError: |
|
DEFAULT_PROTOCOL = pickle.HIGHEST_PROTOCOL |
|
|
|
_SocketType = TypeVar("_SocketType", bound="Socket") |
|
|
|
|
|
class _SocketContext(Generic[_SocketType]): |
|
"""Context Manager for socket bind/unbind""" |
|
|
|
socket: _SocketType |
|
kind: str |
|
addr: str |
|
|
|
def __repr__(self): |
|
return f"<SocketContext({self.kind}={self.addr!r})>" |
|
|
|
def __init__( |
|
self: _SocketContext[_SocketType], socket: _SocketType, kind: str, addr: str |
|
): |
|
assert kind in {"bind", "connect"} |
|
self.socket = socket |
|
self.kind = kind |
|
self.addr = addr |
|
|
|
def __enter__(self: _SocketContext[_SocketType]) -> _SocketType: |
|
return self.socket |
|
|
|
def __exit__(self, *args): |
|
if self.socket.closed: |
|
return |
|
if self.kind == "bind": |
|
self.socket.unbind(self.addr) |
|
elif self.kind == "connect": |
|
self.socket.disconnect(self.addr) |
|
|
|
|
|
SocketReturnType = TypeVar("SocketReturnType") |
|
|
|
|
|
class Socket(SocketBase, AttributeSetter, Generic[SocketReturnType]): |
|
"""The ZMQ socket object |
|
|
|
To create a Socket, first create a Context:: |
|
|
|
ctx = zmq.Context.instance() |
|
|
|
then call ``ctx.socket(socket_type)``:: |
|
|
|
s = ctx.socket(zmq.ROUTER) |
|
|
|
.. versionadded:: 25 |
|
|
|
Sockets can now be shadowed by passing another Socket. |
|
This helps in creating an async copy of a sync socket or vice versa:: |
|
|
|
s = zmq.Socket(async_socket) |
|
|
|
Which previously had to be:: |
|
|
|
s = zmq.Socket.shadow(async_socket.underlying) |
|
""" |
|
|
|
_shadow = False |
|
_shadow_obj = None |
|
_monitor_socket = None |
|
_type_name = 'UNKNOWN' |
|
|
|
@overload |
|
def __init__( |
|
self: Socket[bytes], |
|
ctx_or_socket: zmq.Context, |
|
socket_type: int, |
|
*, |
|
copy_threshold: int | None = None, |
|
): ... |
|
|
|
@overload |
|
def __init__( |
|
self: Socket[bytes], |
|
*, |
|
shadow: Socket | int, |
|
copy_threshold: int | None = None, |
|
): ... |
|
|
|
@overload |
|
def __init__( |
|
self: Socket[bytes], |
|
ctx_or_socket: Socket, |
|
): ... |
|
|
|
def __init__( |
|
self: Socket[bytes], |
|
ctx_or_socket: zmq.Context | Socket | None = None, |
|
socket_type: int = 0, |
|
*, |
|
shadow: Socket | int = 0, |
|
copy_threshold: int | None = None, |
|
): |
|
if isinstance(ctx_or_socket, zmq.Socket): |
|
|
|
shadow = ctx_or_socket |
|
ctx_or_socket = None |
|
|
|
shadow_address: int = 0 |
|
|
|
if shadow: |
|
self._shadow = True |
|
|
|
self._shadow_obj = shadow |
|
if not isinstance(shadow, int): |
|
try: |
|
shadow = cast(int, shadow.underlying) |
|
except AttributeError: |
|
pass |
|
shadow_address = cast_int_addr(shadow) |
|
else: |
|
self._shadow = False |
|
|
|
super().__init__( |
|
ctx_or_socket, |
|
socket_type, |
|
shadow=shadow_address, |
|
copy_threshold=copy_threshold, |
|
) |
|
|
|
try: |
|
socket_type = cast(int, self.get(zmq.TYPE)) |
|
except Exception: |
|
pass |
|
else: |
|
try: |
|
self.__dict__["type"] = stype = SocketType(socket_type) |
|
except ValueError: |
|
self._type_name = str(socket_type) |
|
else: |
|
self._type_name = stype.name |
|
|
|
def __del__(self): |
|
if not self._shadow and not self.closed: |
|
if warn is not None: |
|
|
|
warn( |
|
f"Unclosed socket {self}", |
|
ResourceWarning, |
|
stacklevel=2, |
|
source=self, |
|
) |
|
self.close() |
|
|
|
_repr_cls = "zmq.Socket" |
|
|
|
def __repr__(self): |
|
cls = self.__class__ |
|
|
|
_repr_cls = cls.__dict__.get("_repr_cls", None) |
|
if _repr_cls is None: |
|
_repr_cls = f"{cls.__module__}.{cls.__name__}" |
|
|
|
closed = ' closed' if self._closed else '' |
|
|
|
return f"<{_repr_cls}(zmq.{self._type_name}) at {hex(id(self))}{closed}>" |
|
|
|
|
|
def __enter__(self: _SocketType) -> _SocketType: |
|
"""Sockets are context managers |
|
|
|
.. versionadded:: 14.4 |
|
""" |
|
return self |
|
|
|
def __exit__(self, *args, **kwargs): |
|
self.close() |
|
|
|
|
|
|
|
|
|
|
|
def __copy__(self: _SocketType, memo=None) -> _SocketType: |
|
"""Copying a Socket creates a shadow copy""" |
|
return self.__class__.shadow(self.underlying) |
|
|
|
__deepcopy__ = __copy__ |
|
|
|
@classmethod |
|
def shadow(cls: type[_SocketType], address: int | zmq.Socket) -> _SocketType: |
|
"""Shadow an existing libzmq socket |
|
|
|
address is a zmq.Socket or an integer (or FFI pointer) |
|
representing the address of the libzmq socket. |
|
|
|
.. versionadded:: 14.1 |
|
|
|
.. versionadded:: 25 |
|
Support for shadowing `zmq.Socket` objects, |
|
instead of just integer addresses. |
|
""" |
|
return cls(shadow=address) |
|
|
|
def close(self, linger=None) -> None: |
|
""" |
|
Close the socket. |
|
|
|
If linger is specified, LINGER sockopt will be set prior to closing. |
|
|
|
Note: closing a zmq Socket may not close the underlying sockets |
|
if there are undelivered messages. |
|
Only after all messages are delivered or discarded by reaching the socket's LINGER timeout |
|
(default: forever) |
|
will the underlying sockets be closed. |
|
|
|
This can be called to close the socket by hand. If this is not |
|
called, the socket will automatically be closed when it is |
|
garbage collected, |
|
in which case you may see a ResourceWarning about the unclosed socket. |
|
""" |
|
if self.context: |
|
self.context._rm_socket(self) |
|
super().close(linger=linger) |
|
|
|
|
|
|
|
|
|
|
|
def _connect_cm(self: _SocketType, addr: str) -> _SocketContext[_SocketType]: |
|
"""Context manager to disconnect on exit |
|
|
|
.. versionadded:: 20.0 |
|
""" |
|
return _SocketContext(self, 'connect', addr) |
|
|
|
def _bind_cm(self: _SocketType, addr: str) -> _SocketContext[_SocketType]: |
|
"""Context manager to unbind on exit |
|
|
|
.. versionadded:: 20.0 |
|
""" |
|
try: |
|
|
|
|
|
|
|
addr = cast(bytes, self.get(zmq.LAST_ENDPOINT)).decode("utf8") |
|
except (AttributeError, ZMQError, UnicodeDecodeError): |
|
pass |
|
return _SocketContext(self, 'bind', addr) |
|
|
|
def bind(self: _SocketType, addr: str) -> _SocketContext[_SocketType]: |
|
"""s.bind(addr) |
|
|
|
Bind the socket to an address. |
|
|
|
This causes the socket to listen on a network port. Sockets on the |
|
other side of this connection will use ``Socket.connect(addr)`` to |
|
connect to this socket. |
|
|
|
Returns a context manager which will call unbind on exit. |
|
|
|
.. versionadded:: 20.0 |
|
Can be used as a context manager. |
|
|
|
.. versionadded:: 26.0 |
|
binding to port 0 can be used as a context manager |
|
for binding to a random port. |
|
The URL can be retrieved as `socket.last_endpoint`. |
|
|
|
Parameters |
|
---------- |
|
addr : str |
|
The address string. This has the form 'protocol://interface:port', |
|
for example 'tcp://127.0.0.1:5555'. Protocols supported include |
|
tcp, udp, pgm, epgm, inproc and ipc. If the address is unicode, it is |
|
encoded to utf-8 first. |
|
|
|
""" |
|
try: |
|
super().bind(addr) |
|
except ZMQError as e: |
|
e.strerror += f" (addr={addr!r})" |
|
raise |
|
return self._bind_cm(addr) |
|
|
|
def connect(self: _SocketType, addr: str) -> _SocketContext[_SocketType]: |
|
"""s.connect(addr) |
|
|
|
Connect to a remote 0MQ socket. |
|
|
|
Returns a context manager which will call disconnect on exit. |
|
|
|
.. versionadded:: 20.0 |
|
Can be used as a context manager. |
|
|
|
Parameters |
|
---------- |
|
addr : str |
|
The address string. This has the form 'protocol://interface:port', |
|
for example 'tcp://127.0.0.1:5555'. Protocols supported are |
|
tcp, udp, pgm, inproc and ipc. If the address is unicode, it is |
|
encoded to utf-8 first. |
|
|
|
""" |
|
try: |
|
super().connect(addr) |
|
except ZMQError as e: |
|
e.strerror += f" (addr={addr!r})" |
|
raise |
|
return self._connect_cm(addr) |
|
|
|
|
|
|
|
|
|
|
|
@property |
|
def socket_type(self) -> int: |
|
warn("Socket.socket_type is deprecated, use Socket.type", DeprecationWarning) |
|
return cast(int, self.type) |
|
|
|
|
|
|
|
|
|
|
|
def __dir__(self): |
|
keys = dir(self.__class__) |
|
keys.extend(SocketOption.__members__) |
|
return keys |
|
|
|
|
|
|
|
|
|
setsockopt = SocketBase.set |
|
getsockopt = SocketBase.get |
|
|
|
def __setattr__(self, key, value): |
|
"""Override to allow setting zmq.[UN]SUBSCRIBE even though we have a subscribe method""" |
|
if key in self.__dict__: |
|
object.__setattr__(self, key, value) |
|
return |
|
_key = key.lower() |
|
if _key in ('subscribe', 'unsubscribe'): |
|
if isinstance(value, str): |
|
value = value.encode('utf8') |
|
if _key == 'subscribe': |
|
self.set(zmq.SUBSCRIBE, value) |
|
else: |
|
self.set(zmq.UNSUBSCRIBE, value) |
|
return |
|
super().__setattr__(key, value) |
|
|
|
def fileno(self) -> int: |
|
"""Return edge-triggered file descriptor for this socket. |
|
|
|
This is a read-only edge-triggered file descriptor for both read and write events on this socket. |
|
It is important that all available events be consumed when an event is detected, |
|
otherwise the read event will not trigger again. |
|
|
|
.. versionadded:: 17.0 |
|
""" |
|
return self.FD |
|
|
|
def subscribe(self, topic: str | bytes) -> None: |
|
"""Subscribe to a topic |
|
|
|
Only for SUB sockets. |
|
|
|
.. versionadded:: 15.3 |
|
""" |
|
if isinstance(topic, str): |
|
topic = topic.encode('utf8') |
|
self.set(zmq.SUBSCRIBE, topic) |
|
|
|
def unsubscribe(self, topic: str | bytes) -> None: |
|
"""Unsubscribe from a topic |
|
|
|
Only for SUB sockets. |
|
|
|
.. versionadded:: 15.3 |
|
""" |
|
if isinstance(topic, str): |
|
topic = topic.encode('utf8') |
|
self.set(zmq.UNSUBSCRIBE, topic) |
|
|
|
def set_string(self, option: int, optval: str, encoding='utf-8') -> None: |
|
"""Set socket options with a unicode object. |
|
|
|
This is simply a wrapper for setsockopt to protect from encoding ambiguity. |
|
|
|
See the 0MQ documentation for details on specific options. |
|
|
|
Parameters |
|
---------- |
|
option : int |
|
The name of the option to set. Can be any of: SUBSCRIBE, |
|
UNSUBSCRIBE, IDENTITY |
|
optval : str |
|
The value of the option to set. |
|
encoding : str |
|
The encoding to be used, default is utf8 |
|
""" |
|
if not isinstance(optval, str): |
|
raise TypeError(f"strings only, not {type(optval)}: {optval!r}") |
|
return self.set(option, optval.encode(encoding)) |
|
|
|
setsockopt_unicode = setsockopt_string = set_string |
|
|
|
def get_string(self, option: int, encoding='utf-8') -> str: |
|
"""Get the value of a socket option. |
|
|
|
See the 0MQ documentation for details on specific options. |
|
|
|
Parameters |
|
---------- |
|
option : int |
|
The option to retrieve. |
|
|
|
Returns |
|
------- |
|
optval : str |
|
The value of the option as a unicode string. |
|
""" |
|
if SocketOption(option)._opt_type != _OptType.bytes: |
|
raise TypeError(f"option {option} will not return a string to be decoded") |
|
return cast(bytes, self.get(option)).decode(encoding) |
|
|
|
getsockopt_unicode = getsockopt_string = get_string |
|
|
|
def bind_to_random_port( |
|
self: _SocketType, |
|
addr: str, |
|
min_port: int = 49152, |
|
max_port: int = 65536, |
|
max_tries: int = 100, |
|
) -> int: |
|
"""Bind this socket to a random port in a range. |
|
|
|
If the port range is unspecified, the system will choose the port. |
|
|
|
Parameters |
|
---------- |
|
addr : str |
|
The address string without the port to pass to ``Socket.bind()``. |
|
min_port : int, optional |
|
The minimum port in the range of ports to try (inclusive). |
|
max_port : int, optional |
|
The maximum port in the range of ports to try (exclusive). |
|
max_tries : int, optional |
|
The maximum number of bind attempts to make. |
|
|
|
Returns |
|
------- |
|
port : int |
|
The port the socket was bound to. |
|
|
|
Raises |
|
------ |
|
ZMQBindError |
|
if `max_tries` reached before successful bind |
|
""" |
|
if ( |
|
(zmq.zmq_version_info() >= (3, 2)) |
|
and min_port == 49152 |
|
and max_port == 65536 |
|
): |
|
|
|
|
|
self.bind(f"{addr}:*") |
|
url = cast(bytes, self.last_endpoint).decode('ascii', 'replace') |
|
_, port_s = url.rsplit(':', 1) |
|
return int(port_s) |
|
|
|
for i in range(max_tries): |
|
try: |
|
port = random.randrange(min_port, max_port) |
|
self.bind(f'{addr}:{port}') |
|
except ZMQError as exception: |
|
en = exception.errno |
|
if en == zmq.EADDRINUSE: |
|
continue |
|
elif sys.platform == 'win32' and en == errno.EACCES: |
|
continue |
|
else: |
|
raise |
|
else: |
|
return port |
|
raise ZMQBindError("Could not bind socket to random port.") |
|
|
|
def get_hwm(self) -> int: |
|
"""Get the High Water Mark. |
|
|
|
On libzmq ≥ 3, this gets SNDHWM if available, otherwise RCVHWM |
|
""" |
|
major = zmq.zmq_version_info()[0] |
|
if major >= 3: |
|
|
|
try: |
|
return cast(int, self.get(zmq.SNDHWM)) |
|
except zmq.ZMQError: |
|
pass |
|
|
|
return cast(int, self.get(zmq.RCVHWM)) |
|
else: |
|
return cast(int, self.get(zmq.HWM)) |
|
|
|
def set_hwm(self, value: int) -> None: |
|
"""Set the High Water Mark. |
|
|
|
On libzmq ≥ 3, this sets both SNDHWM and RCVHWM |
|
|
|
|
|
.. warning:: |
|
|
|
New values only take effect for subsequent socket |
|
bind/connects. |
|
""" |
|
major = zmq.zmq_version_info()[0] |
|
if major >= 3: |
|
raised = None |
|
try: |
|
self.sndhwm = value |
|
except Exception as e: |
|
raised = e |
|
try: |
|
self.rcvhwm = value |
|
except Exception as e: |
|
raised = e |
|
|
|
if raised: |
|
raise raised |
|
else: |
|
self.set(zmq.HWM, value) |
|
|
|
hwm = property( |
|
get_hwm, |
|
set_hwm, |
|
None, |
|
"""Property for High Water Mark. |
|
|
|
Setting hwm sets both SNDHWM and RCVHWM as appropriate. |
|
It gets SNDHWM if available, otherwise RCVHWM. |
|
""", |
|
) |
|
|
|
|
|
|
|
|
|
|
|
@overload |
|
def send( |
|
self, |
|
data: Any, |
|
flags: int = ..., |
|
copy: bool = ..., |
|
*, |
|
track: Literal[True], |
|
routing_id: int | None = ..., |
|
group: str | None = ..., |
|
) -> zmq.MessageTracker: ... |
|
|
|
@overload |
|
def send( |
|
self, |
|
data: Any, |
|
flags: int = ..., |
|
copy: bool = ..., |
|
*, |
|
track: Literal[False], |
|
routing_id: int | None = ..., |
|
group: str | None = ..., |
|
) -> None: ... |
|
|
|
@overload |
|
def send( |
|
self, |
|
data: Any, |
|
flags: int = ..., |
|
*, |
|
copy: bool = ..., |
|
routing_id: int | None = ..., |
|
group: str | None = ..., |
|
) -> None: ... |
|
|
|
@overload |
|
def send( |
|
self, |
|
data: Any, |
|
flags: int = ..., |
|
copy: bool = ..., |
|
track: bool = ..., |
|
routing_id: int | None = ..., |
|
group: str | None = ..., |
|
) -> zmq.MessageTracker | None: ... |
|
|
|
def send( |
|
self, |
|
data: Any, |
|
flags: int = 0, |
|
copy: bool = True, |
|
track: bool = False, |
|
routing_id: int | None = None, |
|
group: str | None = None, |
|
) -> zmq.MessageTracker | None: |
|
"""Send a single zmq message frame on this socket. |
|
|
|
This queues the message to be sent by the IO thread at a later time. |
|
|
|
With flags=NOBLOCK, this raises :class:`ZMQError` if the queue is full; |
|
otherwise, this waits until space is available. |
|
See :class:`Poller` for more general non-blocking I/O. |
|
|
|
Parameters |
|
---------- |
|
data : bytes, Frame, memoryview |
|
The content of the message. This can be any object that provides |
|
the Python buffer API (i.e. `memoryview(data)` can be called). |
|
flags : int |
|
0, NOBLOCK, SNDMORE, or NOBLOCK|SNDMORE. |
|
copy : bool |
|
Should the message be sent in a copying or non-copying manner. |
|
track : bool |
|
Should the message be tracked for notification that ZMQ has |
|
finished with it? (ignored if copy=True) |
|
routing_id : int |
|
For use with SERVER sockets |
|
group : str |
|
For use with RADIO sockets |
|
|
|
Returns |
|
------- |
|
None : if `copy` or not track |
|
None if message was sent, raises an exception otherwise. |
|
MessageTracker : if track and not copy |
|
a MessageTracker object, whose `done` property will |
|
be False until the send is completed. |
|
|
|
Raises |
|
------ |
|
TypeError |
|
If a unicode object is passed |
|
ValueError |
|
If `track=True`, but an untracked Frame is passed. |
|
ZMQError |
|
If the send does not succeed for any reason (including |
|
if NOBLOCK is set and the outgoing queue is full). |
|
|
|
|
|
.. versionchanged:: 17.0 |
|
|
|
DRAFT support for routing_id and group arguments. |
|
""" |
|
if routing_id is not None: |
|
if not isinstance(data, zmq.Frame): |
|
data = zmq.Frame( |
|
data, |
|
track=track, |
|
copy=copy or None, |
|
copy_threshold=self.copy_threshold, |
|
) |
|
data.routing_id = routing_id |
|
if group is not None: |
|
if not isinstance(data, zmq.Frame): |
|
data = zmq.Frame( |
|
data, |
|
track=track, |
|
copy=copy or None, |
|
copy_threshold=self.copy_threshold, |
|
) |
|
data.group = group |
|
return super().send(data, flags=flags, copy=copy, track=track) |
|
|
|
def send_multipart( |
|
self, |
|
msg_parts: Sequence, |
|
flags: int = 0, |
|
copy: bool = True, |
|
track: bool = False, |
|
**kwargs, |
|
): |
|
"""Send a sequence of buffers as a multipart message. |
|
|
|
The zmq.SNDMORE flag is added to all msg parts before the last. |
|
|
|
Parameters |
|
---------- |
|
msg_parts : iterable |
|
A sequence of objects to send as a multipart message. Each element |
|
can be any sendable object (Frame, bytes, buffer-providers) |
|
flags : int, optional |
|
Any valid flags for :func:`Socket.send`. |
|
SNDMORE is added automatically for frames before the last. |
|
copy : bool, optional |
|
Should the frame(s) be sent in a copying or non-copying manner. |
|
If copy=False, frames smaller than self.copy_threshold bytes |
|
will be copied anyway. |
|
track : bool, optional |
|
Should the frame(s) be tracked for notification that ZMQ has |
|
finished with it (ignored if copy=True). |
|
|
|
Returns |
|
------- |
|
None : if copy or not track |
|
MessageTracker : if track and not copy |
|
a MessageTracker object, whose `done` property will |
|
be False until the last send is completed. |
|
""" |
|
|
|
for i, msg in enumerate(msg_parts): |
|
if isinstance(msg, (zmq.Frame, bytes, memoryview)): |
|
continue |
|
try: |
|
memoryview(msg) |
|
except Exception: |
|
rmsg = repr(msg) |
|
if len(rmsg) > 32: |
|
rmsg = rmsg[:32] + '...' |
|
raise TypeError( |
|
"Frame %i (%s) does not support the buffer interface." |
|
% ( |
|
i, |
|
rmsg, |
|
) |
|
) |
|
for msg in msg_parts[:-1]: |
|
self.send(msg, zmq.SNDMORE | flags, copy=copy, track=track) |
|
|
|
return self.send(msg_parts[-1], flags, copy=copy, track=track) |
|
|
|
@overload |
|
def recv_multipart( |
|
self, flags: int = ..., *, copy: Literal[True], track: bool = ... |
|
) -> list[bytes]: ... |
|
|
|
@overload |
|
def recv_multipart( |
|
self, flags: int = ..., *, copy: Literal[False], track: bool = ... |
|
) -> list[zmq.Frame]: ... |
|
|
|
@overload |
|
def recv_multipart(self, flags: int = ..., *, track: bool = ...) -> list[bytes]: ... |
|
|
|
@overload |
|
def recv_multipart( |
|
self, flags: int = 0, copy: bool = True, track: bool = False |
|
) -> list[zmq.Frame] | list[bytes]: ... |
|
|
|
def recv_multipart( |
|
self, flags: int = 0, copy: bool = True, track: bool = False |
|
) -> list[zmq.Frame] | list[bytes]: |
|
"""Receive a multipart message as a list of bytes or Frame objects |
|
|
|
Parameters |
|
---------- |
|
flags : int, optional |
|
Any valid flags for :func:`Socket.recv`. |
|
copy : bool, optional |
|
Should the message frame(s) be received in a copying or non-copying manner? |
|
If False a Frame object is returned for each part, if True a copy of |
|
the bytes is made for each frame. |
|
track : bool, optional |
|
Should the message frame(s) be tracked for notification that ZMQ has |
|
finished with it? (ignored if copy=True) |
|
|
|
Returns |
|
------- |
|
msg_parts : list |
|
A list of frames in the multipart message; either Frames or bytes, |
|
depending on `copy`. |
|
|
|
Raises |
|
------ |
|
ZMQError |
|
for any of the reasons :func:`~Socket.recv` might fail |
|
""" |
|
parts = [self.recv(flags, copy=copy, track=track)] |
|
|
|
while self.getsockopt(zmq.RCVMORE): |
|
part = self.recv(flags, copy=copy, track=track) |
|
parts.append(part) |
|
|
|
|
|
return cast(Union[List[zmq.Frame], List[bytes]], parts) |
|
|
|
def _deserialize( |
|
self, |
|
recvd: bytes, |
|
load: Callable[[bytes], Any], |
|
) -> Any: |
|
"""Deserialize a received message |
|
|
|
Override in subclass (e.g. Futures) if recvd is not the raw bytes. |
|
|
|
The default implementation expects bytes and returns the deserialized message immediately. |
|
|
|
Parameters |
|
---------- |
|
|
|
load: callable |
|
Callable that deserializes bytes |
|
recvd: |
|
The object returned by self.recv |
|
|
|
""" |
|
return load(recvd) |
|
|
|
def send_serialized(self, msg, serialize, flags=0, copy=True, **kwargs): |
|
"""Send a message with a custom serialization function. |
|
|
|
.. versionadded:: 17 |
|
|
|
Parameters |
|
---------- |
|
msg : The message to be sent. Can be any object serializable by `serialize`. |
|
serialize : callable |
|
The serialization function to use. |
|
serialize(msg) should return an iterable of sendable message frames |
|
(e.g. bytes objects), which will be passed to send_multipart. |
|
flags : int, optional |
|
Any valid flags for :func:`Socket.send`. |
|
copy : bool, optional |
|
Whether to copy the frames. |
|
|
|
""" |
|
frames = serialize(msg) |
|
return self.send_multipart(frames, flags=flags, copy=copy, **kwargs) |
|
|
|
def recv_serialized(self, deserialize, flags=0, copy=True): |
|
"""Receive a message with a custom deserialization function. |
|
|
|
.. versionadded:: 17 |
|
|
|
Parameters |
|
---------- |
|
deserialize : callable |
|
The deserialization function to use. |
|
deserialize will be called with one argument: the list of frames |
|
returned by recv_multipart() and can return any object. |
|
flags : int, optional |
|
Any valid flags for :func:`Socket.recv`. |
|
copy : bool, optional |
|
Whether to recv bytes or Frame objects. |
|
|
|
Returns |
|
------- |
|
obj : object |
|
The object returned by the deserialization function. |
|
|
|
Raises |
|
------ |
|
ZMQError |
|
for any of the reasons :func:`~Socket.recv` might fail |
|
""" |
|
frames = self.recv_multipart(flags=flags, copy=copy) |
|
return self._deserialize(frames, deserialize) |
|
|
|
def send_string( |
|
self, |
|
u: str, |
|
flags: int = 0, |
|
copy: bool = True, |
|
encoding: str = 'utf-8', |
|
**kwargs, |
|
) -> zmq.Frame | None: |
|
"""Send a Python unicode string as a message with an encoding. |
|
|
|
0MQ communicates with raw bytes, so you must encode/decode |
|
text (str) around 0MQ. |
|
|
|
Parameters |
|
---------- |
|
u : str |
|
The unicode string to send. |
|
flags : int, optional |
|
Any valid flags for :func:`Socket.send`. |
|
encoding : str |
|
The encoding to be used |
|
""" |
|
if not isinstance(u, str): |
|
raise TypeError("str objects only") |
|
return self.send(u.encode(encoding), flags=flags, copy=copy, **kwargs) |
|
|
|
send_unicode = send_string |
|
|
|
def recv_string(self, flags: int = 0, encoding: str = 'utf-8') -> str: |
|
"""Receive a unicode string, as sent by send_string. |
|
|
|
Parameters |
|
---------- |
|
flags : int |
|
Any valid flags for :func:`Socket.recv`. |
|
encoding : str |
|
The encoding to be used |
|
|
|
Returns |
|
------- |
|
s : str |
|
The Python unicode string that arrives as encoded bytes. |
|
|
|
Raises |
|
------ |
|
ZMQError |
|
for any of the reasons :func:`Socket.recv` might fail |
|
""" |
|
msg = self.recv(flags=flags) |
|
return self._deserialize(msg, lambda buf: buf.decode(encoding)) |
|
|
|
recv_unicode = recv_string |
|
|
|
def send_pyobj( |
|
self, obj: Any, flags: int = 0, protocol: int = DEFAULT_PROTOCOL, **kwargs |
|
) -> zmq.Frame | None: |
|
"""Send a Python object as a message using pickle to serialize. |
|
|
|
Parameters |
|
---------- |
|
obj : Python object |
|
The Python object to send. |
|
flags : int |
|
Any valid flags for :func:`Socket.send`. |
|
protocol : int |
|
The pickle protocol number to use. The default is pickle.DEFAULT_PROTOCOL |
|
where defined, and pickle.HIGHEST_PROTOCOL elsewhere. |
|
""" |
|
msg = pickle.dumps(obj, protocol) |
|
return self.send(msg, flags=flags, **kwargs) |
|
|
|
def recv_pyobj(self, flags: int = 0) -> Any: |
|
"""Receive a Python object as a message using pickle to serialize. |
|
|
|
Parameters |
|
---------- |
|
flags : int |
|
Any valid flags for :func:`Socket.recv`. |
|
|
|
Returns |
|
------- |
|
obj : Python object |
|
The Python object that arrives as a message. |
|
|
|
Raises |
|
------ |
|
ZMQError |
|
for any of the reasons :func:`~Socket.recv` might fail |
|
""" |
|
msg = self.recv(flags) |
|
return self._deserialize(msg, pickle.loads) |
|
|
|
def send_json(self, obj: Any, flags: int = 0, **kwargs) -> None: |
|
"""Send a Python object as a message using json to serialize. |
|
|
|
Keyword arguments are passed on to json.dumps |
|
|
|
Parameters |
|
---------- |
|
obj : Python object |
|
The Python object to send |
|
flags : int |
|
Any valid flags for :func:`Socket.send` |
|
""" |
|
send_kwargs = {} |
|
for key in ('routing_id', 'group'): |
|
if key in kwargs: |
|
send_kwargs[key] = kwargs.pop(key) |
|
msg = jsonapi.dumps(obj, **kwargs) |
|
return self.send(msg, flags=flags, **send_kwargs) |
|
|
|
def recv_json(self, flags: int = 0, **kwargs) -> list | str | int | float | dict: |
|
"""Receive a Python object as a message using json to serialize. |
|
|
|
Keyword arguments are passed on to json.loads |
|
|
|
Parameters |
|
---------- |
|
flags : int |
|
Any valid flags for :func:`Socket.recv`. |
|
|
|
Returns |
|
------- |
|
obj : Python object |
|
The Python object that arrives as a message. |
|
|
|
Raises |
|
------ |
|
ZMQError |
|
for any of the reasons :func:`~Socket.recv` might fail |
|
""" |
|
msg = self.recv(flags) |
|
return self._deserialize(msg, lambda buf: jsonapi.loads(buf, **kwargs)) |
|
|
|
_poller_class = Poller |
|
|
|
def poll(self, timeout: int | None = None, flags: int = zmq.POLLIN) -> int: |
|
"""Poll the socket for events. |
|
|
|
See :class:`Poller` to wait for multiple sockets at once. |
|
|
|
Parameters |
|
---------- |
|
timeout : int |
|
The timeout (in milliseconds) to wait for an event. If unspecified |
|
(or specified None), will wait forever for an event. |
|
flags : int |
|
default: POLLIN. |
|
POLLIN, POLLOUT, or POLLIN|POLLOUT. The event flags to poll for. |
|
|
|
Returns |
|
------- |
|
event_mask : int |
|
The poll event mask (POLLIN, POLLOUT), |
|
0 if the timeout was reached without an event. |
|
""" |
|
|
|
if self.closed: |
|
raise ZMQError(zmq.ENOTSUP) |
|
|
|
p = self._poller_class() |
|
p.register(self, flags) |
|
evts = dict(p.poll(timeout)) |
|
|
|
return evts.get(self, 0) |
|
|
|
def get_monitor_socket( |
|
self: _SocketType, events: int | None = None, addr: str | None = None |
|
) -> _SocketType: |
|
"""Return a connected PAIR socket ready to receive the event notifications. |
|
|
|
.. versionadded:: libzmq-4.0 |
|
.. versionadded:: 14.0 |
|
|
|
Parameters |
|
---------- |
|
events : int |
|
default: `zmq.EVENT_ALL` |
|
The bitmask defining which events are wanted. |
|
addr : str |
|
The optional endpoint for the monitoring sockets. |
|
|
|
Returns |
|
------- |
|
socket : zmq.Socket |
|
The PAIR socket, connected and ready to receive messages. |
|
""" |
|
|
|
if zmq.zmq_version_info() < (4,): |
|
raise NotImplementedError( |
|
f"get_monitor_socket requires libzmq >= 4, have {zmq.zmq_version()}" |
|
) |
|
|
|
|
|
if self._monitor_socket: |
|
if self._monitor_socket.closed: |
|
self._monitor_socket = None |
|
else: |
|
return self._monitor_socket |
|
|
|
if addr is None: |
|
|
|
addr = f"inproc://monitor.s-{self.FD}" |
|
if events is None: |
|
|
|
events = zmq.EVENT_ALL |
|
|
|
self.monitor(addr, events) |
|
|
|
self._monitor_socket = self.context.socket(zmq.PAIR) |
|
self._monitor_socket.connect(addr) |
|
return self._monitor_socket |
|
|
|
def disable_monitor(self) -> None: |
|
"""Shutdown the PAIR socket (created using get_monitor_socket) |
|
that is serving socket events. |
|
|
|
.. versionadded:: 14.4 |
|
""" |
|
self._monitor_socket = None |
|
self.monitor(None, 0) |
|
|
|
|
|
SyncSocket: TypeAlias = Socket[bytes] |
|
|
|
__all__ = ['Socket', 'SyncSocket'] |
|
|