diff --git a/examples/mongodb/controller.py b/examples/mongodb/controller.py index 68a40905c..9177b51af 100644 --- a/examples/mongodb/controller.py +++ b/examples/mongodb/controller.py @@ -71,7 +71,7 @@ def start(self) -> None: if len(msg) != 3: error_msg = f'invalid message received: {msg}' print(error_msg) - reply = [msg[0], error_msg] + reply = [msg[0], error_msg.encode()] socket.send_multipart(reply) continue id = msg[0] @@ -81,11 +81,11 @@ def start(self) -> None: reply = [id] if operation == 'add': self.add_document(contents) - reply.append("success") + reply.append(b"success") elif operation == 'get': doc = self.get_document_by_keys(contents) json_doc = self._doc_to_json(doc) - reply.append(json_doc) + reply.append(json_doc.encode()) else: print('unknown request') socket.send_multipart(reply) diff --git a/examples/poll/pair.py b/examples/poll/pair.py index db0933e39..f135e5334 100644 --- a/examples/poll/pair.py +++ b/examples/poll/pair.py @@ -34,8 +34,8 @@ assert socks[s2] == zmq.POLLOUT # Now do a send on both, wait and test for zmq.POLLOUT|zmq.POLLIN -s1.send('msg1') -s2.send('msg2') +s1.send(b'msg1') +s2.send(b'msg2') time.sleep(1.0) socks = dict(poller.poll()) assert socks[s1] == zmq.POLLOUT | zmq.POLLIN diff --git a/examples/poll/pubsub.py b/examples/poll/pubsub.py index 94952a190..84a010370 100644 --- a/examples/poll/pubsub.py +++ b/examples/poll/pubsub.py @@ -17,7 +17,7 @@ ctx = zmq.Context() s1 = ctx.socket(zmq.PUB) s2 = ctx.socket(zmq.SUB) -s2.setsockopt(zmq.SUBSCRIBE, '') +s2.setsockopt(zmq.SUBSCRIBE, b'') s1.bind(addr) s2.connect(addr) @@ -35,7 +35,7 @@ assert s2 not in socks # Make sure that s1 stays in POLLOUT after a send. -s1.send('msg1') +s1.send(b'msg1') socks = dict(poller.poll()) assert socks[s1] == zmq.POLLOUT diff --git a/examples/poll/reqrep.py b/examples/poll/reqrep.py index d6a458c4c..afb8d5671 100644 --- a/examples/poll/reqrep.py +++ b/examples/poll/reqrep.py @@ -34,7 +34,7 @@ assert socks[s2] == zmq.POLLOUT # Make sure that s2 goes immediately into state 0 after send. -s2.send('msg1') +s2.send(b'msg1') socks = dict(poller.poll()) assert s2 not in socks @@ -49,7 +49,7 @@ assert socks[s1] == zmq.POLLOUT # Make sure s1 goes into state 0 after send. -s1.send('msg2') +s1.send(b'msg2') socks = dict(poller.poll()) assert s1 not in socks diff --git a/examples/win32-interrupt/display.py b/examples/win32-interrupt/display.py index bcff213b2..ef4d722c6 100644 --- a/examples/win32-interrupt/display.py +++ b/examples/win32-interrupt/display.py @@ -11,7 +11,7 @@ def main(addrs: list[str]): control = context.socket(zmq.PUB) control.bind('inproc://control') updates = context.socket(zmq.SUB) - updates.setsockopt(zmq.SUBSCRIBE, "") + updates.setsockopt(zmq.SUBSCRIBE, b"") updates.connect('inproc://control') for addr in addrs: print("Connecting to: ", addr) diff --git a/zmq/backend/__init__.pyi b/zmq/backend/__init__.pyi index 5ed569eec..3d0611508 100644 --- a/zmq/backend/__init__.pyi +++ b/zmq/backend/__init__.pyi @@ -1,10 +1,12 @@ -from typing import Any, Callable, Optional, TypeVar, Union, overload +from collections.abc import Sequence +from typing import Any, Final, Protocol, overload, type_check_only -from typing_extensions import Literal +from _typeshed import HasFileno +from typing_extensions import Buffer, Literal, Self import zmq -from .select import select_backend +from .select import select_backend as select_backend __all__ = [ 'Context', @@ -24,31 +26,49 @@ __all__ = [ 'PYZMQ_DRAFT_API', ] -# avoid collision in Frame.bytes -_bytestr = bytes - -T = TypeVar("T") +IPC_PATH_MAX_LEN: Final[int] = ... +PYZMQ_DRAFT_API: Final[bool] = ... class Frame: - buffer: Any - bytes: bytes more: bool tracker: Any + def __init__( self, data: Any = None, track: bool = False, copy: bool | None = None, copy_threshold: int | None = None, - ): ... - def copy_fast(self: T) -> T: ... - def get(self, option: int) -> int | _bytestr | str: ... - def set(self, option: int, value: int | _bytestr | str) -> None: ... + ) -> None: ... + def __len__(self) -> int: ... + def __copy__(self) -> Self: ... + def fast_copy(self) -> Self: ... + + # + @overload + def get(self, option: int | Literal["routing_id"]) -> int: ... + @overload + def get(self, option: bytes | Literal["group"]) -> str: ... + @overload + def get(self, option: str) -> int | str: ... + + # + @overload + def set(self, option: Literal["routing_id"], value: int) -> None: ... + @overload + def set(self, option: Literal["group"], value: str | bytes) -> None: ... + @overload + def set(self, option: int, value: int) -> None: ... + + # + @property + def buffer(self) -> memoryview: ... + @property + def bytes(self) -> bytes: ... Message = Frame class Socket: - underlying: int context: zmq.Context copy_threshold: int @@ -58,86 +78,162 @@ class Socket: def __init__( self, context: Context | None = None, - socket_type: int = 0, + socket_type: int = -1, shadow: int = 0, - copy_threshold: int | None = zmq.COPY_THRESHOLD, + copy_threshold: int | None = None, + ) -> None: ... + + # + @property + def underlying(self) -> int: ... + @property + def closed(self) -> bool: ... + + # + def close(self, linger: int | None = None) -> None: ... + + # + def set(self, option: int, optval: int | bytes) -> None: ... + def get(self, option: int) -> int | bytes: ... + + # + def bind(self, addr: str) -> None: ... + def connect(self, addr: str) -> None: ... + def unbind(self, addr: str | bytes) -> None: ... + def disconnect(self, addr: str | bytes) -> None: ... + def monitor( + self, + addr: str | bytes | None, + events: int = zmq.EVENT_ALL, ) -> None: ... - def close(self, linger: int | None = ...) -> None: ... - def get(self, option: int) -> int | bytes | str: ... - def set(self, option: int, value: int | bytes | str) -> None: ... - def connect(self, url: str): ... - def disconnect(self, url: str) -> None: ... - def bind(self, url: str): ... - def unbind(self, url: str) -> None: ... + def join(self, group: str | bytes) -> None: ... + def leave(self, group: str | bytes) -> None: ... + + # + @overload # data: Buffer, copy=True (default) def send( self, - data: Any, - flags: int = ..., - copy: bool = ..., - track: bool = ..., - ) -> zmq.MessageTracker | None: ... - @overload - def recv( + data: Buffer, + flags: int = 0, + copy: Literal[True] = True, + track: bool = False, + ) -> None: ... + @overload # data: Buffer, copy=False (keyword) + def send( self, - flags: int = ..., + data: Buffer, + flags: int = 0, *, copy: Literal[False], - track: bool = ..., - ) -> zmq.Frame: ... - @overload + track: bool = False, + ) -> zmq.MessageTracker: ... + @overload # data: Buffer, copy=False (positional) + def send( + self, + data: Buffer, + flags: int, + copy: Literal[False], + track: bool = False, + ) -> zmq.MessageTracker: ... + @overload # data: Frame + def send( + self, + data: Frame, + flags: int = 0, + copy: bool = True, + track: bool = False, + ) -> zmq.MessageTracker: ... + + # + @overload # copy=True (default) def recv( self, - flags: int = ..., - *, - copy: Literal[True], - track: bool = ..., + flags: int = 0, + copy: Literal[True] = True, + track: bool = False, ) -> bytes: ... - @overload + @overload # copy=False (keyword) + def recv( + self, + flags: int = 0, + *, + copy: Literal[False], + track: bool = False, + ) -> zmq.Frame: ... + @overload # copy=False (positional) def recv( self, - flags: int = ..., + flags: int, + copy: Literal[False], track: bool = False, - ) -> bytes: ... - @overload + ) -> zmq.Frame: ... + @overload # fallback overload (mypy bug workaround) def recv( self, - flags: int | None = ..., - copy: bool = ..., - track: bool | None = False, - ) -> zmq.Frame | bytes: ... - def recv_into(self, buf, /, *, nbytes: int = 0, flags: int = 0) -> int: ... - def monitor(self, addr: str | None, events: int) -> None: ... - # draft methods - def join(self, group: str) -> None: ... - def leave(self, group: str) -> None: ... + flags: int = 0, + copy: bool = True, + track: bool = False, + ) -> bytes | zmq.Frame: ... + + # + def recv_into( + self, buffer: Buffer, /, *, nbytes: int = 0, flags: int = 0 + ) -> int: ... class Context: - underlying: int - def __init__(self, io_threads: int = 1, shadow: int = 0): ... - def get(self, option: int) -> int | bytes | str: ... - def set(self, option: int, value: int | bytes | str) -> None: ... - def socket(self, socket_type: int) -> Socket: ... - def term(self) -> None: ... + handle: int + closed: bool -IPC_PATH_MAX_LEN: int -PYZMQ_DRAFT_API: bool + def __init__(self, io_threads: int = 1, shadow: int = 0) -> None: ... + + # + @property + def underlying(self) -> int: ... + + # + def term(self) -> None: ... + def set(self, option: int, optval: int) -> None: ... + def get(self, option: int) -> int: ... -def has(capability: str) -> bool: ... -def curve_keypair() -> tuple[bytes, bytes]: ... -def curve_public(secret_key: bytes) -> bytes: ... -def strerror(errno: int | None = ...) -> str: ... def zmq_errno() -> int: ... -def zmq_version() -> str: ... +def strerror(errno: int) -> str: ... def zmq_version_info() -> tuple[int, int, int]: ... +def has(capability: str) -> bool: ... +def curve_keypair() -> tuple[bytes, bytes]: ... +def curve_public(secret_key: str | bytes) -> bytes: ... + +# +@overload def zmq_poll( - sockets: list[Any], timeout: int | None = ... + sockets: Sequence[tuple[Socket, int]], + timeout: int = -1, ) -> list[tuple[Socket, int]]: ... +@overload +def zmq_poll( + sockets: Sequence[tuple[int | HasFileno, int]], + timeout: int = -1, +) -> list[tuple[int, int]]: ... + +# def proxy(frontend: Socket, backend: Socket, capture: Socket | None = None) -> int: ... def proxy_steerable( frontend: Socket, backend: Socket, - capture: Socket | None = ..., - control: Socket | None = ..., + capture: Socket | None = None, + control: Socket | None = None, ) -> int: ... -monitored_queue = Callable | None +@type_check_only +class _MonitoredQueueFunction(Protocol): + def __call__( + self, + /, + in_socket: Socket, + out_socket: Socket, + mon_socket: Socket, + in_prefix: Buffer = b"in", + out_prefix: Buffer = b"out", + ) -> int: ... + +# None for the cffi backend +monitored_queue: Final[_MonitoredQueueFunction | None] = ... diff --git a/zmq/backend/select.py b/zmq/backend/select.py index ecf89b04e..b5f4a2f02 100644 --- a/zmq/backend/select.py +++ b/zmq/backend/select.py @@ -4,6 +4,7 @@ # Distributed under the terms of the Modified BSD License. from importlib import import_module +from typing import Any public_api = [ 'Context', @@ -24,7 +25,7 @@ ] -def select_backend(name: str) -> dict: +def select_backend(name: str) -> dict[str, Any]: """Select the pyzmq backend""" try: mod = import_module(name) diff --git a/zmq/sugar/attrsettr.py b/zmq/sugar/attrsettr.py index 844fce606..2663a8477 100644 --- a/zmq/sugar/attrsettr.py +++ b/zmq/sugar/attrsettr.py @@ -5,7 +5,7 @@ from __future__ import annotations import errno -from typing import TypeVar, Union +from typing import Any, TypeVar, Union from .. import constants @@ -14,7 +14,7 @@ class AttributeSetter: - def __setattr__(self, key: str, value: OptValT) -> None: + def __setattr__(self, key: str, value: Any) -> None: """set zmq options by attribute""" if key in self.__dict__: @@ -36,7 +36,7 @@ def __setattr__(self, key: str, value: OptValT) -> None: else: self._set_attr_opt(upper_key, opt, value) - def _set_attr_opt(self, name: str, opt: int, value: OptValT) -> None: + def _set_attr_opt(self, name: str, opt: int, value: Any) -> None: """override if setattr should do something other than call self.set""" self.set(opt, value) @@ -67,11 +67,11 @@ def _get_attr_opt(self, name, opt) -> OptValT: """override if getattr should do something other than call self.get""" return self.get(opt) - def get(self, opt: int) -> OptValT: + def get(self, option: int) -> OptValT: """Override in subclass""" raise NotImplementedError("override in subclass") - def set(self, opt: int, val: OptValT) -> None: + def set(self, option: int, value: Any) -> None: """Override in subclass""" raise NotImplementedError("override in subclass") diff --git a/zmq/sugar/context.py b/zmq/sugar/context.py index 21bc2cdda..e16bf0a32 100644 --- a/zmq/sugar/context.py +++ b/zmq/sugar/context.py @@ -8,7 +8,7 @@ import atexit import os from threading import Lock -from typing import Any, Callable, Generic, TypeVar, overload +from typing import Any, Callable, Generic, TypeVar, cast, overload from warnings import warn from weakref import WeakSet @@ -381,7 +381,7 @@ def getsockopt(self, opt: int) -> OptValT: def _set_attr_opt(self, name: str, opt: int, value: OptValT) -> None: """set default sockopts as attributes""" if name in ContextOption.__members__: - return self.set(opt, value) + return self.set(opt, cast(int, value)) elif name in SocketOption.__members__: self.sockopts[opt] = value else: diff --git a/zmq/sugar/frame.py b/zmq/sugar/frame.py index 3239d357e..4c1b87d4d 100644 --- a/zmq/sugar/frame.py +++ b/zmq/sugar/frame.py @@ -3,13 +3,20 @@ # Copyright (C) PyZMQ Developers # Distributed under the terms of the Modified BSD License. +from __future__ import annotations + +from typing import Literal, overload + import zmq from zmq.backend import Frame as FrameBase from .attrsettr import AttributeSetter -def _draft(v, feature): +def _draft( + v: tuple[int] | tuple[int, int] | tuple[int, int, int], + feature: str, +) -> None: zmq.error._check_version(v, feature) if not zmq.DRAFT_API: raise RuntimeError( @@ -65,11 +72,17 @@ class Frame(FrameBase, AttributeSetter): will be copied and messages larger than this will be shared with libzmq. """ - def __getitem__(self, key): + @overload + def __getitem__(self, key: int | Literal["routing_id"]) -> int: ... + @overload + def __getitem__(self, key: bytes | Literal["group"]) -> str: ... + @overload + def __getitem__(self, key: str) -> int | str: ... + def __getitem__(self, key: int | str | bytes) -> int | str: # map Frame['User-Id'] to Frame.get('User-Id') return self.get(key) - def __repr__(self): + def __repr__(self) -> str: """Return the str form of the message.""" nbytes = len(self) msg_suffix = "" @@ -97,7 +110,7 @@ def __repr__(self): return f"<{_module}.{self.__class__.__name__}({msg_bytes!r}{msg_suffix})>" @property - def group(self): + def group(self) -> str: """The RADIO-DISH group of the message. Requires libzmq >= 4.2 and pyzmq built with draft APIs enabled. @@ -108,12 +121,12 @@ def group(self): return self.get('group') @group.setter - def group(self, group): + def group(self, group: str | bytes) -> None: _draft((4, 2), "RADIO-DISH") self.set('group', group) @property - def routing_id(self): + def routing_id(self) -> int: """The CLIENT-SERVER routing id of the message. Requires libzmq >= 4.2 and pyzmq built with draft APIs enabled. @@ -124,11 +137,12 @@ def routing_id(self): return self.get('routing_id') @routing_id.setter - def routing_id(self, routing_id): + def routing_id(self, routing_id: int) -> None: _draft((4, 2), "CLIENT-SERVER") self.set('routing_id', routing_id) # keep deprecated alias Message = Frame + __all__ = ['Frame', 'Message'] diff --git a/zmq/sugar/socket.py b/zmq/sugar/socket.py index dac805e92..fbbc2fb3f 100644 --- a/zmq/sugar/socket.py +++ b/zmq/sugar/socket.py @@ -11,12 +11,12 @@ import sys from collections.abc import Sequence from typing import ( + TYPE_CHECKING, Any, Callable, Generic, Literal, TypeVar, - Union, cast, overload, ) @@ -33,37 +33,39 @@ from .attrsettr import AttributeSetter from .poll import Poller +if TYPE_CHECKING: + from typing_extensions import Buffer, Self + try: DEFAULT_PROTOCOL = pickle.DEFAULT_PROTOCOL except AttributeError: DEFAULT_PROTOCOL = pickle.HIGHEST_PROTOCOL -_SocketType = TypeVar("_SocketType", bound="Socket") +_T = TypeVar("_T") +_RT = TypeVar("_RT") +_SocketT_co = TypeVar("_SocketT_co", bound="Socket", covariant=True) -_JSONType: TypeAlias = ( - "int | float | str | bool | list[_JSONType] | dict[str, _JSONType]" -) +# must match the `jsonapi.loads` return type +_JSON: TypeAlias = "dict[str, Any] | list[Any] | str | float" -class _SocketContext(Generic[_SocketType]): +class _SocketContext(Generic[_SocketT_co]): """Context Manager for socket bind/unbind""" - socket: _SocketType + socket: _SocketT_co kind: str addr: str def __repr__(self): return f"" - def __init__( - self: _SocketContext[_SocketType], socket: _SocketType, kind: str, addr: str - ): + def __init__(self, socket: _SocketT_co, kind: str, addr: str): assert kind in {"bind", "connect"} self.socket = socket self.kind = kind self.addr = addr - def __enter__(self: _SocketContext[_SocketType]) -> _SocketType: + def __enter__(self) -> _SocketT_co: return self.socket def __exit__(self, *args): @@ -75,10 +77,10 @@ def __exit__(self, *args): self.socket.disconnect(self.addr) -SocketReturnType = TypeVar("SocketReturnType") +_SocketReturnT_co = TypeVar("_SocketReturnT_co", covariant=True) -class Socket(SocketBase, AttributeSetter, Generic[SocketReturnType]): +class Socket(SocketBase, AttributeSetter, Generic[_SocketReturnT_co]): """The ZMQ socket object To create a Socket, first create a Context:: @@ -106,6 +108,8 @@ class Socket(SocketBase, AttributeSetter, Generic[SocketReturnType]): _monitor_socket = None _type_name = 'UNKNOWN' + context: zmq.Context + @overload def __init__( self: Socket[bytes], @@ -113,22 +117,19 @@ def __init__( socket_type: int, *, copy_threshold: int | None = None, - ): ... - + ) -> None: ... @overload def __init__( self: Socket[bytes], *, shadow: Socket | int, copy_threshold: int | None = None, - ): ... - + ) -> None: ... @overload def __init__( self: Socket[bytes], ctx_or_socket: Socket, - ): ... - + ) -> None: ... def __init__( self: Socket[bytes], ctx_or_socket: zmq.Context | Socket | None = None, @@ -136,7 +137,7 @@ def __init__( *, shadow: Socket | int = 0, copy_threshold: int | None = None, - ): + ) -> None: shadow_context: zmq.Context | None = None if isinstance(ctx_or_socket, zmq.Socket): # positional Socket(other_socket) @@ -182,7 +183,7 @@ def __init__( else: self._type_name = stype.name - def __del__(self): + def __del__(self) -> None: if not self._shadow and not self.closed: if warn is not None: # warn can be None during process teardown @@ -196,7 +197,7 @@ def __del__(self): _repr_cls = "zmq.Socket" - def __repr__(self): + def __repr__(self) -> str: cls = self.__class__ # look up _repr_cls on exact class, not inherited _repr_cls = cls.__dict__.get("_repr_cls", None) @@ -208,28 +209,28 @@ def __repr__(self): return f"<{_repr_cls}(zmq.{self._type_name}) at {hex(id(self))}{closed}>" # socket as context manager: - def __enter__(self: _SocketType) -> _SocketType: + def __enter__(self) -> Self: """Sockets are context managers .. versionadded:: 14.4 """ return self - def __exit__(self, *args, **kwargs): + def __exit__(self, *args: Any, **kwargs: Any) -> None: self.close() # ------------------------------------------------------------------------- # Socket creation # ------------------------------------------------------------------------- - def __copy__(self: _SocketType, memo=None) -> _SocketType: + def __copy__(self, memo: Any | None = None) -> Self: """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: + def shadow(cls, address: int | zmq.Socket) -> Self: """Shadow an existing libzmq socket address is a zmq.Socket or an integer (or FFI pointer) @@ -243,7 +244,7 @@ def shadow(cls: type[_SocketType], address: int | zmq.Socket) -> _SocketType: """ return cls(shadow=address) - def close(self, linger=None) -> None: + def close(self, linger: int | None = None) -> None: """ Close the socket. @@ -268,14 +269,14 @@ def close(self, linger=None) -> None: # Connect/Bind context managers # ------------------------------------------------------------------------- - def _connect_cm(self: _SocketType, addr: str) -> _SocketContext[_SocketType]: + def _connect_cm(self, addr: str) -> _SocketContext[Self]: """Context manager to disconnect on exit .. versionadded:: 20.0 """ return _SocketContext(self, 'connect', addr) - def _bind_cm(self: _SocketType, addr: str) -> _SocketContext[_SocketType]: + def _bind_cm(self, addr: str) -> _SocketContext[Self]: """Context manager to unbind on exit .. versionadded:: 20.0 @@ -289,7 +290,7 @@ def _bind_cm(self: _SocketType, addr: str) -> _SocketContext[_SocketType]: pass return _SocketContext(self, 'bind', addr) - def bind(self: _SocketType, addr: str) -> _SocketContext[_SocketType]: + def bind(self, addr: str) -> _SocketContext[Self]: # type:ignore[override] """s.bind(addr) Bind the socket to an address. @@ -324,7 +325,7 @@ def bind(self: _SocketType, addr: str) -> _SocketContext[_SocketType]: raise return self._bind_cm(addr) - def connect(self: _SocketType, addr: str) -> _SocketContext[_SocketType]: + def connect(self, addr: str) -> _SocketContext[Self]: # type:ignore[override] """s.connect(addr) Connect to a remote 0MQ socket. @@ -363,7 +364,7 @@ def socket_type(self) -> int: # Hooks for sockopt completion # ------------------------------------------------------------------------- - def __dir__(self): + def __dir__(self) -> list[str]: keys = dir(self.__class__) keys.extend(SocketOption.__members__) return keys @@ -374,7 +375,7 @@ def __dir__(self): setsockopt = SocketBase.set getsockopt = SocketBase.get - def __setattr__(self, key, value): + def __setattr__(self, key: str, value: Any) -> None: """Override to allow setting zmq.[UN]SUBSCRIBE even though we have a subscribe method""" if key in self.__dict__: object.__setattr__(self, key, value) @@ -423,7 +424,7 @@ def unsubscribe(self, topic: str | bytes) -> None: topic = topic.encode('utf8') self.set(zmq.UNSUBSCRIBE, topic) - def set_string(self, option: int, optval: str, encoding='utf-8') -> None: + def set_string(self, option: int, optval: str, encoding: str = 'utf-8') -> None: """Set socket options with a unicode object. This is simply a wrapper for setsockopt to protect from encoding ambiguity. @@ -446,7 +447,7 @@ def set_string(self, option: int, optval: str, encoding='utf-8') -> None: setsockopt_unicode = setsockopt_string = set_string - def get_string(self, option: int, encoding='utf-8') -> str: + def get_string(self, option: int, encoding: str = 'utf-8') -> str: """Get the value of a socket option. See the 0MQ documentation for details on specific options. @@ -468,7 +469,7 @@ def get_string(self, option: int, encoding='utf-8') -> str: getsockopt_unicode = getsockopt_string = get_string def bind_to_random_port( - self: _SocketType, + self, addr: str, min_port: int = 49152, max_port: int = 65536, @@ -560,67 +561,80 @@ def set_hwm(self, value: int) -> None: if raised: raise raised - hwm = property( - get_hwm, - set_hwm, - None, - """Property for High Water Mark. + if TYPE_CHECKING: - Setting hwm sets both SNDHWM and RCVHWM as appropriate. - It gets SNDHWM if available, otherwise RCVHWM. - """, - ) + @property + def hwm(self) -> int: ... + @hwm.setter + def hwm(self, hwm: int, /) -> None: ... + + else: + hwm: property = 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. + """, + ) # ------------------------------------------------------------------------- # Sending and receiving messages # ------------------------------------------------------------------------- - @overload + @overload # type:ignore[override] # data: Buffer, copy=True (default) def send( self, - data: Any, - flags: int = ..., - copy: bool = ..., - *, - track: Literal[True], - routing_id: int | None = ..., - group: str | None = ..., - ) -> zmq.MessageTracker: ... - - @overload + data: Buffer, + flags: int = 0, + copy: Literal[True] = True, + track: bool = False, + routing_id: int | None = None, + group: str | None = None, + ) -> None: ... + @overload # data: Buffer, copy=False (keyword) def send( self, - data: Any, - flags: int = ..., - copy: bool = ..., + data: Buffer, + flags: int = 0, *, - track: Literal[False], - routing_id: int | None = ..., - group: str | None = ..., - ) -> None: ... - - @overload + copy: Literal[False], + track: bool = False, + routing_id: int | None = None, + group: str | None = None, + ) -> zmq.MessageTracker: ... + @overload # data: Buffer, copy=False (positional) def send( self, - data: Any, - flags: int = ..., - *, - copy: bool = ..., - routing_id: int | None = ..., - group: str | None = ..., - ) -> None: ... - - @overload + data: Buffer, + flags: int, + copy: Literal[False], + track: bool = False, + routing_id: int | None = None, + group: str | None = None, + ) -> zmq.MessageTracker: ... + @overload # data: Buffer, copy=True|False (mypy bug workaround) def send( self, - data: Any, - flags: int = ..., - copy: bool = ..., - track: bool = ..., - routing_id: int | None = ..., - group: str | None = ..., + data: Buffer, + flags: int, + copy: bool = True, + track: bool = False, + routing_id: int | None = None, + group: str | None = None, ) -> zmq.MessageTracker | None: ... - + @overload # data: Frame + def send( + self, + data: zmq.Frame, + flags: int = 0, + copy: bool = True, + track: bool = False, + routing_id: int | None = None, + group: str | None = None, + ) -> zmq.MessageTracker: ... def send( self, data: Any, @@ -700,12 +714,12 @@ def send( def send_multipart( self, - msg_parts: Sequence, + msg_parts: Sequence[zmq.Frame | Buffer], flags: int = 0, copy: bool = True, track: bool = False, - **kwargs, - ): + **kwargs: Any, + ) -> zmq.MessageTracker | None: """Send a sequence of buffers as a multipart message. The zmq.SNDMORE flag is added to all msg parts before the last. @@ -751,27 +765,25 @@ def send_multipart( # Send the last part without the extra SNDMORE flag. return self.send(msg_parts[-1], flags, copy=copy, track=track) - @overload + @overload # copy=True (default) def recv_multipart( - self, flags: int = ..., *, copy: Literal[True], track: bool = ... + self, flags: int = 0, copy: Literal[True] = True, track: bool = False ) -> list[bytes]: ... - - @overload + @overload # copy=False (keyword) def recv_multipart( - self, flags: int = ..., *, copy: Literal[False], track: bool = ... + self, flags: int = 0, *, copy: Literal[False], track: bool = False ) -> list[zmq.Frame]: ... - - @overload - def recv_multipart(self, flags: int = ..., *, track: bool = ...) -> list[bytes]: ... - - @overload + @overload # copy=False (positional) def recv_multipart( - self, flags: int = 0, copy: bool = True, track: bool = False - ) -> list[zmq.Frame] | list[bytes]: ... - + self, flags: int, copy: Literal[False], track: bool = False + ) -> list[zmq.Frame]: ... + @overload # (mypy bug workaround) + def recv_multipart( + self, flags: int, copy: bool, track: bool = False + ) -> list[bytes] | list[zmq.Frame]: ... def recv_multipart( self, flags: int = 0, copy: bool = True, track: bool = False - ) -> list[zmq.Frame] | list[bytes]: + ) -> list[Any]: """Receive a multipart message as a list of bytes or Frame objects Parameters @@ -802,15 +814,13 @@ def recv_multipart( while self.getsockopt(zmq.RCVMORE): part = self.recv(flags, copy=copy, track=track) parts.append(part) - # cast List[Union] to Union[List] - # how do we get mypy to recognize that return type is invariant on `copy`? - return cast(Union[list[zmq.Frame], list[bytes]], parts) + return parts def _deserialize( self, - recvd: bytes, - load: Callable[[bytes], Any], - ) -> Any: + recvd: _T, + load: Callable[[_T], _RT], + ) -> _RT: """Deserialize a received message Override in subclass (e.g. Futures) if recvd is not the raw bytes. @@ -828,7 +838,14 @@ def _deserialize( """ return load(recvd) - def send_serialized(self, msg, serialize, flags=0, copy=True, **kwargs): + def send_serialized( + self, + msg: _T, + serialize: Callable[[_T], Sequence[zmq.Frame | Buffer]], + flags: int = 0, + copy: bool = True, + **kwargs: Any, + ) -> zmq.MessageTracker | None: """Send a message with a custom serialization function. .. versionadded:: 17 @@ -849,7 +866,28 @@ def send_serialized(self, msg, serialize, flags=0, copy=True, **kwargs): frames = serialize(msg) return self.send_multipart(frames, flags=flags, copy=copy, **kwargs) - def recv_serialized(self, deserialize, flags=0, copy=True): + # these overloads should not be needed, but mypy seems to have trouble inferring + # the type of deserialize's argument without them + @overload + def recv_serialized( + self, + deserialize: Callable[[list[bytes]], _RT], + flags: int = 0, + copy: bool = True, + ) -> _RT: ... + @overload + def recv_serialized( + self, + deserialize: Callable[[list[zmq.Frame]], _RT], + flags: int = 0, + copy: bool = True, + ) -> _RT: ... + def recv_serialized( + self, + deserialize: Callable[[list[Any]], _RT], + flags: int = 0, + copy: bool = True, + ) -> _RT: """Receive a message with a custom deserialization function. .. versionadded:: 17 @@ -884,8 +922,8 @@ def send_string( flags: int = 0, copy: bool = True, encoding: str = 'utf-8', - **kwargs, - ) -> zmq.Frame | None: + **kwargs: Any, + ) -> zmq.MessageTracker | None: """Send a Python unicode string as a message with an encoding. 0MQ communicates with raw bytes, so you must encode/decode @@ -932,8 +970,12 @@ def recv_string(self, flags: int = 0, encoding: str = 'utf-8') -> str: recv_unicode = recv_string def send_pyobj( - self, obj: Any, flags: int = 0, protocol: int = DEFAULT_PROTOCOL, **kwargs - ) -> zmq.Frame | None: + self, + obj: object, + flags: int = 0, + protocol: int = DEFAULT_PROTOCOL, + **kwargs: Any, + ) -> zmq.MessageTracker | None: """ Send a Python object as a message using pickle to serialize. @@ -990,7 +1032,7 @@ def recv_pyobj(self, flags: int = 0) -> Any: msg = self.recv(flags) return self._deserialize(msg, pickle.loads) - def send_json(self, obj: Any, flags: int = 0, **kwargs) -> None: + def send_json(self, obj: object, flags: int = 0, **kwargs: Any) -> None: """Send a Python object as a message using json to serialize. Keyword arguments are passed on to json.dumps @@ -1009,7 +1051,7 @@ def send_json(self, obj: Any, flags: int = 0, **kwargs) -> None: msg = jsonapi.dumps(obj, **kwargs) return self.send(msg, flags=flags, **send_kwargs) - def recv_json(self, flags: int = 0, **kwargs) -> _JSONType: + def recv_json(self, flags: int = 0, **kwargs: Any) -> _JSON: """Receive a Python object as a message using json to serialize. Keyword arguments are passed on to json.loads @@ -1065,8 +1107,8 @@ def poll(self, timeout: int | None = None, flags: int = zmq.POLLIN) -> int: return evts.get(self, 0) def get_monitor_socket( - self: _SocketType, events: int | None = None, addr: str | None = None - ) -> _SocketType: + self, events: int | None = None, addr: str | None = None + ) -> Self: """Return a connected PAIR socket ready to receive the event notifications. .. versionadded:: libzmq-4.0