File size: 3,310 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
"""Module holding utility and convenience functions for zmq event monitoring."""

# Copyright (C) PyZMQ Developers
# Distributed under the terms of the Modified BSD License.

from __future__ import annotations

import struct
from typing import Awaitable, overload

import zmq
import zmq.asyncio
from zmq._typing import TypedDict
from zmq.error import _check_version


class _MonitorMessage(TypedDict):
    event: int
    value: int
    endpoint: bytes


def parse_monitor_message(msg: list[bytes]) -> _MonitorMessage:
    """decode zmq_monitor event messages.

    Parameters
    ----------
    msg : list(bytes)
        zmq multipart message that has arrived on a monitor PAIR socket.

        First frame is::

            16 bit event id
            32 bit event value
            no padding

        Second frame is the endpoint as a bytestring

    Returns
    -------
    event : dict
        event description as dict with the keys `event`, `value`, and `endpoint`.
    """
    if len(msg) != 2 or len(msg[0]) != 6:
        raise RuntimeError(f"Invalid event message format: {msg}")
    event_id, value = struct.unpack("=hi", msg[0])
    event: _MonitorMessage = {
        'event': zmq.Event(event_id),
        'value': zmq.Event(value),
        'endpoint': msg[1],
    }
    return event


async def _parse_monitor_msg_async(
    awaitable_msg: Awaitable[list[bytes]],
) -> _MonitorMessage:
    """Like parse_monitor_msg, but awaitable

    Given awaitable message, return awaitable for the parsed monitor message.
    """

    msg = await awaitable_msg
    # 4.0-style event API
    return parse_monitor_message(msg)


@overload
def recv_monitor_message(
    socket: zmq.asyncio.Socket,
    flags: int = 0,
) -> Awaitable[_MonitorMessage]: ...


@overload
def recv_monitor_message(
    socket: zmq.Socket[bytes],
    flags: int = 0,
) -> _MonitorMessage: ...


def recv_monitor_message(
    socket: zmq.Socket,
    flags: int = 0,
) -> _MonitorMessage | Awaitable[_MonitorMessage]:
    """Receive and decode the given raw message from the monitoring socket and return a dict.

    Requires libzmq ≥ 4.0

    The returned dict will have the following entries:
      event : int
        the event id as described in `libzmq.zmq_socket_monitor`
      value : int
        the event value associated with the event, see `libzmq.zmq_socket_monitor`
      endpoint : str
        the affected endpoint

    .. versionchanged:: 23.1
        Support for async sockets added.
        When called with a async socket,
        returns an awaitable for the monitor message.

    Parameters
    ----------
    socket : zmq.Socket
        The PAIR socket (created by other.get_monitor_socket()) on which to recv the message
    flags : int
        standard zmq recv flags

    Returns
    -------
    event : dict
        event description as dict with the keys `event`, `value`, and `endpoint`.
    """

    _check_version((4, 0), 'libzmq event API')
    # will always return a list
    msg = socket.recv_multipart(flags)

    # transparently handle asyncio socket,
    # returns a Future instead of a dict
    if isinstance(msg, Awaitable):
        return _parse_monitor_msg_async(msg)

    # 4.0-style event API
    return parse_monitor_message(msg)


__all__ = ['parse_monitor_message', 'recv_monitor_message']