You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
421 lines
14 KiB
Python
421 lines
14 KiB
Python
"""Python bindings for 0MQ."""
|
|
|
|
# Copyright (C) PyZMQ Developers
|
|
# Distributed under the terms of the Modified BSD License.
|
|
|
|
from __future__ import annotations
|
|
|
|
import atexit
|
|
import os
|
|
from threading import Lock
|
|
from typing import Any, Callable, Generic, TypeVar, overload
|
|
from warnings import warn
|
|
from weakref import WeakSet
|
|
|
|
import zmq
|
|
from zmq._typing import TypeAlias
|
|
from zmq.backend import Context as ContextBase
|
|
from zmq.constants import ContextOption, Errno, SocketOption
|
|
from zmq.error import ZMQError
|
|
from zmq.utils.interop import cast_int_addr
|
|
|
|
from .attrsettr import AttributeSetter, OptValT
|
|
from .socket import Socket, SyncSocket
|
|
|
|
# notice when exiting, to avoid triggering term on exit
|
|
_exiting = False
|
|
|
|
|
|
def _notice_atexit() -> None:
|
|
global _exiting
|
|
_exiting = True
|
|
|
|
|
|
atexit.register(_notice_atexit)
|
|
|
|
_ContextType = TypeVar('_ContextType', bound='Context')
|
|
_SocketType = TypeVar('_SocketType', bound='Socket', covariant=True)
|
|
|
|
|
|
class Context(ContextBase, AttributeSetter, Generic[_SocketType]):
|
|
"""Create a zmq Context
|
|
|
|
A zmq Context creates sockets via its ``ctx.socket`` method.
|
|
|
|
.. versionchanged:: 24
|
|
|
|
When using a Context as a context manager (``with zmq.Context()``),
|
|
or deleting a context without closing it first,
|
|
``ctx.destroy()`` is called,
|
|
closing any leftover sockets,
|
|
instead of `ctx.term()` which requires sockets to be closed first.
|
|
|
|
This prevents hangs caused by `ctx.term()` if sockets are left open,
|
|
but means that unclean destruction of contexts
|
|
(with sockets left open) is not safe
|
|
if sockets are managed in other threads.
|
|
|
|
.. versionadded:: 25
|
|
|
|
Contexts can now be shadowed by passing another Context.
|
|
This helps in creating an async copy of a sync context or vice versa::
|
|
|
|
ctx = zmq.Context(async_ctx)
|
|
|
|
Which previously had to be::
|
|
|
|
ctx = zmq.Context.shadow(async_ctx.underlying)
|
|
"""
|
|
|
|
sockopts: dict[int, Any]
|
|
_instance: Any = None
|
|
_instance_lock = Lock()
|
|
_instance_pid: int | None = None
|
|
_shadow = False
|
|
_shadow_obj = None
|
|
_warn_destroy_close = False
|
|
_sockets: WeakSet
|
|
# mypy doesn't like a default value here
|
|
_socket_class: type[_SocketType] = Socket # type: ignore
|
|
|
|
@overload
|
|
def __init__(self: SyncContext, io_threads: int = 1): ...
|
|
|
|
@overload
|
|
def __init__(self: SyncContext, io_threads: Context, /): ...
|
|
|
|
@overload
|
|
def __init__(self: SyncContext, *, shadow: Context | int): ...
|
|
|
|
def __init__(
|
|
self: SyncContext,
|
|
io_threads: int | Context = 1,
|
|
shadow: Context | int = 0,
|
|
) -> None:
|
|
if isinstance(io_threads, Context):
|
|
# allow positional shadow `zmq.Context(zmq.asyncio.Context())`
|
|
# this s
|
|
shadow = io_threads
|
|
io_threads = 1
|
|
|
|
shadow_address: int = 0
|
|
if shadow:
|
|
self._shadow = True
|
|
# hold a reference to the shadow object
|
|
self._shadow_obj = shadow
|
|
if not isinstance(shadow, int):
|
|
try:
|
|
shadow = shadow.underlying
|
|
except AttributeError:
|
|
pass
|
|
shadow_address = cast_int_addr(shadow)
|
|
else:
|
|
self._shadow = False
|
|
super().__init__(io_threads=io_threads, shadow=shadow_address)
|
|
self.sockopts = {}
|
|
self._sockets = WeakSet()
|
|
|
|
def __del__(self) -> None:
|
|
"""Deleting a Context without closing it destroys it and all sockets.
|
|
|
|
.. versionchanged:: 24
|
|
Switch from threadsafe `term()` which hangs in the event of open sockets
|
|
to less safe `destroy()` which
|
|
warns about any leftover sockets and closes them.
|
|
"""
|
|
|
|
# Calling locals() here conceals issue #1167 on Windows CPython 3.5.4.
|
|
locals()
|
|
|
|
if not self._shadow and not _exiting and not self.closed:
|
|
self._warn_destroy_close = True
|
|
if warn is not None and getattr(self, "_sockets", None) is not None:
|
|
# warn can be None during process teardown
|
|
warn(
|
|
f"Unclosed context {self}",
|
|
ResourceWarning,
|
|
stacklevel=2,
|
|
source=self,
|
|
)
|
|
self.destroy()
|
|
|
|
_repr_cls = "zmq.Context"
|
|
|
|
def __repr__(self) -> str:
|
|
cls = self.__class__
|
|
# look up _repr_cls on exact class, not inherited
|
|
_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 ''
|
|
if getattr(self, "_sockets", None):
|
|
n_sockets = len(self._sockets)
|
|
s = 's' if n_sockets > 1 else ''
|
|
sockets = f"{n_sockets} socket{s}"
|
|
else:
|
|
sockets = ""
|
|
return f"<{_repr_cls}({sockets}) at {hex(id(self))}{closed}>"
|
|
|
|
def __enter__(self: _ContextType) -> _ContextType:
|
|
return self
|
|
|
|
def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
|
|
# warn about any leftover sockets before closing them
|
|
self._warn_destroy_close = True
|
|
self.destroy()
|
|
|
|
def __copy__(self: _ContextType, memo: Any = None) -> _ContextType:
|
|
"""Copying a Context creates a shadow copy"""
|
|
return self.__class__.shadow(self.underlying)
|
|
|
|
__deepcopy__ = __copy__
|
|
|
|
@classmethod
|
|
def shadow(cls: type[_ContextType], address: int | zmq.Context) -> _ContextType:
|
|
"""Shadow an existing libzmq context
|
|
|
|
address is a zmq.Context or an integer (or FFI pointer)
|
|
representing the address of the libzmq context.
|
|
|
|
.. versionadded:: 14.1
|
|
|
|
.. versionadded:: 25
|
|
Support for shadowing `zmq.Context` objects,
|
|
instead of just integer addresses.
|
|
"""
|
|
return cls(shadow=address)
|
|
|
|
@classmethod
|
|
def shadow_pyczmq(cls: type[_ContextType], ctx: Any) -> _ContextType:
|
|
"""Shadow an existing pyczmq context
|
|
|
|
ctx is the FFI `zctx_t *` pointer
|
|
|
|
.. versionadded:: 14.1
|
|
"""
|
|
from pyczmq import zctx # type: ignore
|
|
|
|
from zmq.utils.interop import cast_int_addr
|
|
|
|
underlying = zctx.underlying(ctx)
|
|
address = cast_int_addr(underlying)
|
|
return cls(shadow=address)
|
|
|
|
# static method copied from tornado IOLoop.instance
|
|
@classmethod
|
|
def instance(cls: type[_ContextType], io_threads: int = 1) -> _ContextType:
|
|
"""Returns a global Context instance.
|
|
|
|
Most single-process applications have a single, global Context.
|
|
Use this method instead of passing around Context instances
|
|
throughout your code.
|
|
|
|
A common pattern for classes that depend on Contexts is to use
|
|
a default argument to enable programs with multiple Contexts
|
|
but not require the argument for simpler applications::
|
|
|
|
class MyClass(object):
|
|
def __init__(self, context=None):
|
|
self.context = context or Context.instance()
|
|
|
|
.. versionchanged:: 18.1
|
|
|
|
When called in a subprocess after forking,
|
|
a new global instance is created instead of inheriting
|
|
a Context that won't work from the parent process.
|
|
"""
|
|
if (
|
|
cls._instance is None
|
|
or cls._instance_pid != os.getpid()
|
|
or cls._instance.closed
|
|
):
|
|
with cls._instance_lock:
|
|
if (
|
|
cls._instance is None
|
|
or cls._instance_pid != os.getpid()
|
|
or cls._instance.closed
|
|
):
|
|
cls._instance = cls(io_threads=io_threads)
|
|
cls._instance_pid = os.getpid()
|
|
return cls._instance
|
|
|
|
def term(self) -> None:
|
|
"""Close or terminate the context.
|
|
|
|
Context termination is performed in the following steps:
|
|
|
|
- Any blocking operations currently in progress on sockets open within context shall
|
|
raise :class:`zmq.ContextTerminated`.
|
|
With the exception of socket.close(), any further operations on sockets open within this context
|
|
shall raise :class:`zmq.ContextTerminated`.
|
|
- After interrupting all blocking calls, term shall block until the following conditions are satisfied:
|
|
- All sockets open within context have been closed.
|
|
- For each socket within context, all messages sent on the socket have either been
|
|
physically transferred to a network peer,
|
|
or the socket's linger period set with the zmq.LINGER socket option has expired.
|
|
|
|
For further details regarding socket linger behaviour refer to libzmq documentation for ZMQ_LINGER.
|
|
|
|
This can be called to close the context by hand. If this is not called,
|
|
the context will automatically be closed when it is garbage collected,
|
|
in which case you may see a ResourceWarning about the unclosed context.
|
|
"""
|
|
super().term()
|
|
|
|
# -------------------------------------------------------------------------
|
|
# Hooks for ctxopt completion
|
|
# -------------------------------------------------------------------------
|
|
|
|
def __dir__(self) -> list[str]:
|
|
keys = dir(self.__class__)
|
|
keys.extend(ContextOption.__members__)
|
|
return keys
|
|
|
|
# -------------------------------------------------------------------------
|
|
# Creating Sockets
|
|
# -------------------------------------------------------------------------
|
|
|
|
def _add_socket(self, socket: Any) -> None:
|
|
"""Add a weakref to a socket for Context.destroy / reference counting"""
|
|
self._sockets.add(socket)
|
|
|
|
def _rm_socket(self, socket: Any) -> None:
|
|
"""Remove a socket for Context.destroy / reference counting"""
|
|
# allow _sockets to be None in case of process teardown
|
|
if getattr(self, "_sockets", None) is not None:
|
|
self._sockets.discard(socket)
|
|
|
|
def destroy(self, linger: int | None = None) -> None:
|
|
"""Close all sockets associated with this context and then terminate
|
|
the context.
|
|
|
|
.. warning::
|
|
|
|
destroy involves calling :meth:`Socket.close`, which is **NOT** threadsafe.
|
|
If there are active sockets in other threads, this must not be called.
|
|
|
|
Parameters
|
|
----------
|
|
|
|
linger : int, optional
|
|
If specified, set LINGER on sockets prior to closing them.
|
|
"""
|
|
if self.closed:
|
|
return
|
|
|
|
sockets: list[_SocketType] = list(getattr(self, "_sockets", None) or [])
|
|
for s in sockets:
|
|
if s and not s.closed:
|
|
if self._warn_destroy_close and warn is not None:
|
|
# warn can be None during process teardown
|
|
warn(
|
|
f"Destroying context with unclosed socket {s}",
|
|
ResourceWarning,
|
|
stacklevel=3,
|
|
source=s,
|
|
)
|
|
if linger is not None:
|
|
s.setsockopt(SocketOption.LINGER, linger)
|
|
s.close()
|
|
|
|
self.term()
|
|
|
|
def socket(
|
|
self: _ContextType,
|
|
socket_type: int,
|
|
socket_class: Callable[[_ContextType, int], _SocketType] | None = None,
|
|
**kwargs: Any,
|
|
) -> _SocketType:
|
|
"""Create a Socket associated with this Context.
|
|
|
|
Parameters
|
|
----------
|
|
socket_type : int
|
|
The socket type, which can be any of the 0MQ socket types:
|
|
REQ, REP, PUB, SUB, PAIR, DEALER, ROUTER, PULL, PUSH, etc.
|
|
|
|
socket_class: zmq.Socket
|
|
The socket class to instantiate, if different from the default for this Context.
|
|
e.g. for creating an asyncio socket attached to a default Context or vice versa.
|
|
|
|
.. versionadded:: 25
|
|
|
|
kwargs:
|
|
will be passed to the __init__ method of the socket class.
|
|
"""
|
|
if self.closed:
|
|
raise ZMQError(Errno.ENOTSUP)
|
|
if socket_class is None:
|
|
socket_class = self._socket_class
|
|
s: _SocketType = (
|
|
socket_class( # set PYTHONTRACEMALLOC=2 to get the calling frame
|
|
self, socket_type, **kwargs
|
|
)
|
|
)
|
|
for opt, value in self.sockopts.items():
|
|
try:
|
|
s.setsockopt(opt, value)
|
|
except ZMQError:
|
|
# ignore ZMQErrors, which are likely for socket options
|
|
# that do not apply to a particular socket type, e.g.
|
|
# SUBSCRIBE for non-SUB sockets.
|
|
pass
|
|
self._add_socket(s)
|
|
return s
|
|
|
|
def setsockopt(self, opt: int, value: Any) -> None:
|
|
"""set default socket options for new sockets created by this Context
|
|
|
|
.. versionadded:: 13.0
|
|
"""
|
|
self.sockopts[opt] = value
|
|
|
|
def getsockopt(self, opt: int) -> OptValT:
|
|
"""get default socket options for new sockets created by this Context
|
|
|
|
.. versionadded:: 13.0
|
|
"""
|
|
return self.sockopts[opt]
|
|
|
|
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)
|
|
elif name in SocketOption.__members__:
|
|
self.sockopts[opt] = value
|
|
else:
|
|
raise AttributeError(f"No such context or socket option: {name}")
|
|
|
|
def _get_attr_opt(self, name: str, opt: int) -> OptValT:
|
|
"""get default sockopts as attributes"""
|
|
if name in ContextOption.__members__:
|
|
return self.get(opt)
|
|
else:
|
|
if opt not in self.sockopts:
|
|
raise AttributeError(name)
|
|
else:
|
|
return self.sockopts[opt]
|
|
|
|
def __delattr__(self, key: str) -> None:
|
|
"""delete default sockopts as attributes"""
|
|
if key in self.__dict__:
|
|
self.__dict__.pop(key)
|
|
return
|
|
key = key.upper()
|
|
try:
|
|
opt = getattr(SocketOption, key)
|
|
except AttributeError:
|
|
raise AttributeError(f"No such socket option: {key!r}")
|
|
else:
|
|
if opt not in self.sockopts:
|
|
raise AttributeError(key)
|
|
else:
|
|
del self.sockopts[opt]
|
|
|
|
|
|
SyncContext: TypeAlias = Context[SyncSocket]
|
|
|
|
|
|
__all__ = ['Context', 'SyncContext']
|