Skip to content

Commit d95ca07

Browse files
jorenhamjorenham
committed
Typing fixes and improvements for Socket, Frame, and Context 
1 parent 9f811c7 commit d95ca07

6 files changed

Lines changed: 348 additions & 195 deletions

File tree

‎zmq/backend/__init__.pyi‎

Lines changed: 164 additions & 68 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,12 @@
1-
from typing import Any, Callable, Optional, TypeVar, Union, overload
1+
from collections.abc import Sequence
2+
from typing import Any, Final, Protocol, overload, type_check_only
23

3-
from typing_extensions import Literal
4+
from _typeshed import HasFileno
5+
from typing_extensions import Buffer, Literal, Self
46

57
import zmq
68

7-
from .select import select_backend
9+
from .select import select_backend as select_backend
810

911
__all__ = [
1012
'Context',
@@ -24,31 +26,49 @@ __all__ = [
2426
'PYZMQ_DRAFT_API',
2527
]
2628

27-
# avoid collision in Frame.bytes
28-
_bytestr = bytes
29-
30-
T = TypeVar("T")
29+
IPC_PATH_MAX_LEN: Final[int] = ...
30+
PYZMQ_DRAFT_API: Final[bool] = ...
3131

3232
class Frame:
33-
buffer: Any
34-
bytes: bytes
3533
more: bool
3634
tracker: Any
35+
3736
def __init__(
3837
self,
3938
data: Any = None,
4039
track: bool = False,
4140
copy: bool | None = None,
4241
copy_threshold: int | None = None,
43-
): ...
44-
def copy_fast(self: T) -> T: ...
45-
def get(self, option: int) -> int | _bytestr | str: ...
46-
def set(self, option: int, value: int | _bytestr | str) -> None: ...
42+
) -> None: ...
43+
def __len__(self) -> int: ...
44+
def __copy__(self) -> Self: ...
45+
def fast_copy(self) -> Self: ...
46+
47+
#
48+
@overload
49+
def get(self, option: int | Literal["routing_id"]) -> int: ...
50+
@overload
51+
def get(self, option: bytes | Literal["group"]) -> str: ...
52+
@overload
53+
def get(self, option: str) -> int | str: ...
54+
55+
#
56+
@overload
57+
def set(self, option: Literal["routing_id"], value: int) -> None: ...
58+
@overload
59+
def set(self, option: Literal["group"], value: str | bytes) -> None: ...
60+
@overload
61+
def set(self, option: int, value: int) -> None: ...
62+
63+
#
64+
@property
65+
def buffer(self) -> memoryview: ...
66+
@property
67+
def bytes(self) -> bytes: ...
4768

4869
Message = Frame
4970

5071
class Socket:
51-
underlying: int
5272
context: zmq.Context
5373
copy_threshold: int
5474

@@ -58,86 +78,162 @@ class Socket:
5878
def __init__(
5979
self,
6080
context: Context | None = None,
61-
socket_type: int = 0,
81+
socket_type: int = -1,
6282
shadow: int = 0,
63-
copy_threshold: int | None = zmq.COPY_THRESHOLD,
83+
copy_threshold: int | None = None,
84+
) -> None: ...
85+
86+
#
87+
@property
88+
def underlying(self) -> int: ...
89+
@property
90+
def closed(self) -> bool: ...
91+
92+
#
93+
def close(self, linger: int | None = None) -> None: ...
94+
95+
#
96+
def set(self, option: int, optval: int | bytes) -> None: ...
97+
def get(self, option: int) -> int | bytes: ...
98+
99+
#
100+
def bind(self, addr: str) -> None: ...
101+
def connect(self, addr: str) -> None: ...
102+
def unbind(self, addr: str | bytes) -> None: ...
103+
def disconnect(self, addr: str | bytes) -> None: ...
104+
def monitor(
105+
self,
106+
addr: str | bytes | None,
107+
events: int = zmq.EVENT_ALL,
64108
) -> None: ...
65-
def close(self, linger: int | None = ...) -> None: ...
66-
def get(self, option: int) -> int | bytes | str: ...
67-
def set(self, option: int, value: int | bytes | str) -> None: ...
68-
def connect(self, url: str): ...
69-
def disconnect(self, url: str) -> None: ...
70-
def bind(self, url: str): ...
71-
def unbind(self, url: str) -> None: ...
109+
def join(self, group: str | bytes) -> None: ...
110+
def leave(self, group: str | bytes) -> None: ...
111+
112+
#
113+
@overload # data: Buffer, copy=True (default)
72114
def send(
73115
self,
74-
data: Any,
75-
flags: int = ...,
76-
copy: bool = ...,
77-
track: bool = ...,
78-
) -> zmq.MessageTracker | None: ...
79-
@overload
80-
def recv(
116+
data: Buffer,
117+
flags: int = 0,
118+
copy: Literal[True] = True,
119+
track: bool = False,
120+
) -> None: ...
121+
@overload # data: Buffer, copy=False (keyword)
122+
def send(
81123
self,
82-
flags: int = ...,
124+
data: Buffer,
125+
flags: int = 0,
83126
*,
84127
copy: Literal[False],
85-
track: bool = ...,
86-
) -> zmq.Frame: ...
87-
@overload
128+
track: bool = False,
129+
) -> zmq.MessageTracker: ...
130+
@overload # data: Buffer, copy=False (positional)
131+
def send(
132+
self,
133+
data: Buffer,
134+
flags: int,
135+
copy: Literal[False],
136+
track: bool = False,
137+
) -> zmq.MessageTracker: ...
138+
@overload # data: Frame
139+
def send(
140+
self,
141+
data: Frame,
142+
flags: int = 0,
143+
copy: bool = True,
144+
track: bool = False,
145+
) -> zmq.MessageTracker: ...
146+
147+
#
148+
@overload # copy=True (default)
88149
def recv(
89150
self,
90-
flags: int = ...,
91-
*,
92-
copy: Literal[True],
93-
track: bool = ...,
151+
flags: int = 0,
152+
copy: Literal[True] = True,
153+
track: bool = False,
94154
) -> bytes: ...
95-
@overload
155+
@overload # copy=False (keyword)
156+
def recv(
157+
self,
158+
flags: int = 0,
159+
*,
160+
copy: Literal[False],
161+
track: bool = False,
162+
) -> zmq.Frame: ...
163+
@overload # copy=False (positional)
96164
def recv(
97165
self,
98-
flags: int = ...,
166+
flags: int,
167+
copy: Literal[False],
99168
track: bool = False,
100-
) -> bytes: ...
101-
@overload
169+
) -> zmq.Frame: ...
170+
@overload # fallback overload (mypy bug workaround)
102171
def recv(
103172
self,
104-
flags: int | None = ...,
105-
copy: bool = ...,
106-
track: bool | None = False,
107-
) -> zmq.Frame | bytes: ...
108-
def recv_into(self, buf, /, *, nbytes: int = 0, flags: int = 0) -> int: ...
109-
def monitor(self, addr: str | None, events: int) -> None: ...
110-
# draft methods
111-
def join(self, group: str) -> None: ...
112-
def leave(self, group: str) -> None: ...
173+
flags: int = 0,
174+
copy: bool = True,
175+
track: bool = False,
176+
) -> bytes | zmq.Frame: ...
177+
178+
#
179+
def recv_into(
180+
self, buffer: Buffer, /, *, nbytes: int = 0, flags: int = 0
181+
) -> int: ...
113182

114183
class Context:
115-
underlying: int
116-
def __init__(self, io_threads: int = 1, shadow: int = 0): ...
117-
def get(self, option: int) -> int | bytes | str: ...
118-
def set(self, option: int, value: int | bytes | str) -> None: ...
119-
def socket(self, socket_type: int) -> Socket: ...
120-
def term(self) -> None: ...
184+
handle: int
185+
closed: bool
121186

122-
IPC_PATH_MAX_LEN: int
123-
PYZMQ_DRAFT_API: bool
187+
def __init__(self, io_threads: int = 1, shadow: int = 0) -> None: ...
188+
189+
#
190+
@property
191+
def underlying(self) -> int: ...
192+
193+
#
194+
def term(self) -> None: ...
195+
def set(self, option: int, optval: int) -> None: ...
196+
def get(self, option: int) -> int: ...
124197

125-
def has(capability: str) -> bool: ...
126-
def curve_keypair() -> tuple[bytes, bytes]: ...
127-
def curve_public(secret_key: bytes) -> bytes: ...
128-
def strerror(errno: int | None = ...) -> str: ...
129198
def zmq_errno() -> int: ...
130-
def zmq_version() -> str: ...
199+
def strerror(errno: int) -> str: ...
131200
def zmq_version_info() -> tuple[int, int, int]: ...
201+
def has(capability: str) -> bool: ...
202+
def curve_keypair() -> tuple[bytes, bytes]: ...
203+
def curve_public(secret_key: str | bytes) -> bytes: ...
204+
205+
#
206+
@overload
132207
def zmq_poll(
133-
sockets: list[Any], timeout: int | None = ...
208+
sockets: Sequence[tuple[Socket, int]],
209+
timeout: int = -1,
134210
) -> list[tuple[Socket, int]]: ...
211+
@overload
212+
def zmq_poll(
213+
sockets: Sequence[tuple[int | HasFileno, int]],
214+
timeout: int = -1,
215+
) -> list[tuple[int, int]]: ...
216+
217+
#
135218
def proxy(frontend: Socket, backend: Socket, capture: Socket | None = None) -> int: ...
136219
def proxy_steerable(
137220
frontend: Socket,
138221
backend: Socket,
139-
capture: Socket | None = ...,
140-
control: Socket | None = ...,
222+
capture: Socket | None = None,
223+
control: Socket | None = None,
141224
) -> int: ...
142225

143-
monitored_queue = Callable | None
226+
@type_check_only
227+
class _MonitoredQueueFunction(Protocol):
228+
def __call__(
229+
self,
230+
/,
231+
in_socket: Socket,
232+
out_socket: Socket,
233+
mon_socket: Socket,
234+
in_prefix: Buffer = b"in",
235+
out_prefix: Buffer = b"out",
236+
) -> int: ...
237+
238+
# None for the cffi backend
239+
monitored_queue: Final[_MonitoredQueueFunction | None] = ...

‎zmq/backend/select.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
# Distributed under the terms of the Modified BSD License.
55

66
from importlib import import_module
7+
from typing import Any
78

89
public_api = [
910
'Context',
@@ -24,7 +25,7 @@
2425
]
2526

2627

27-
def select_backend(name: str) -> dict:
28+
def select_backend(name: str) -> dict[str, Any]:
2829
"""Select the pyzmq backend"""
2930
try:
3031
mod = import_module(name)

‎zmq/sugar/attrsettr.py‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
from __future__ import annotations
66

77
import errno
8-
from typing import TypeVar, Union
8+
from typing import Any, TypeVar, Union
99

1010
from .. import constants
1111

@@ -14,7 +14,7 @@
1414

1515

1616
class AttributeSetter:
17-
def __setattr__(self, key: str, value: OptValT) -> None:
17+
def __setattr__(self, key: str, value: Any) -> None:
1818
"""set zmq options by attribute"""
1919

2020
if key in self.__dict__:
@@ -36,7 +36,7 @@ def __setattr__(self, key: str, value: OptValT) -> None:
3636
else:
3737
self._set_attr_opt(upper_key, opt, value)
3838

39-
def _set_attr_opt(self, name: str, opt: int, value: OptValT) -> None:
39+
def _set_attr_opt(self, name: str, opt: int, value: Any) -> None:
4040
"""override if setattr should do something other than call self.set"""
4141
self.set(opt, value)
4242

@@ -67,11 +67,11 @@ def _get_attr_opt(self, name, opt) -> OptValT:
6767
"""override if getattr should do something other than call self.get"""
6868
return self.get(opt)
6969

70-
def get(self, opt: int) -> OptValT:
70+
def get(self, option: int) -> OptValT:
7171
"""Override in subclass"""
7272
raise NotImplementedError("override in subclass")
7373

74-
def set(self, opt: int, val: OptValT) -> None:
74+
def set(self, option: int, value: Any) -> None:
7575
"""Override in subclass"""
7676
raise NotImplementedError("override in subclass")
7777

‎zmq/sugar/context.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88
import atexit
99
import os
1010
from threading import Lock
11-
from typing import Any, Callable, Generic, TypeVar, overload
11+
from typing import Any, Callable, Generic, TypeVar, cast, overload
1212
from warnings import warn
1313
from weakref import WeakSet
1414

@@ -381,7 +381,7 @@ def getsockopt(self, opt: int) -> OptValT:
381381
def _set_attr_opt(self, name: str, opt: int, value: OptValT) -> None:
382382
"""set default sockopts as attributes"""
383383
if name in ContextOption.__members__:
384-
return self.set(opt, value)
384+
return self.set(opt, cast(int, value))
385385
elif name in SocketOption.__members__:
386386
self.sockopts[opt] = value
387387
else:

0 commit comments

Comments
 (0)