cpu-diffusers.upgrade-tmp/venv/lib/python3.10/site-packages/zmq/sugar/tracker.py

"""Tracker for zero-copy messages with 0MQ."""

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

from __future__ import annotations

import time
from threading import Event

from zmq.backend import Frame
from zmq.error import NotDone


class MessageTracker:
    """A class for tracking if 0MQ is done using one or more messages.

    When you send a 0MQ message, it is not sent immediately. The 0MQ IO thread
    sends the message at some later time. Often you want to know when 0MQ has
    actually sent the message though. This is complicated by the fact that
    a single 0MQ message can be sent multiple times using different sockets.
    This class allows you to track all of the 0MQ usages of a message.

    Parameters
    ----------
    towatch : Event, MessageTracker, zmq.Frame
        This objects to track. This class can track the low-level
        Events used by the Message class, other MessageTrackers or
        actual Messages.
    """

    events: set[Event]
    peers: set[MessageTracker]

    def __init__(self, *towatch: tuple[MessageTracker | Event | Frame]):
        """Create a message tracker to track a set of messages.

        Parameters
        ----------
        *towatch : tuple of Event, MessageTracker, Message instances.
            This list of objects to track. This class can track the low-level
            Events used by the Message class, other MessageTrackers or
            actual Messages.
        """
        self.events = set()
        self.peers = set()
        for obj in towatch:
            if isinstance(obj, Event):
                self.events.add(obj)
            elif isinstance(obj, MessageTracker):
                self.peers.add(obj)
            elif isinstance(obj, Frame):
                if not obj.tracker:
                    raise ValueError("Not a tracked message")
                self.peers.add(obj.tracker)
            else:
                raise TypeError(f"Require Events or Message Frames, not {type(obj)}")

    @property
    def done(self):
        """Is 0MQ completely done with the message(s) being tracked?"""
        for evt in self.events:
            if not evt.is_set():
                return False
        for pm in self.peers:
            if not pm.done:
                return False
        return True

    def wait(self, timeout: float | int = -1):
        """Wait for 0MQ to be done with the message or until `timeout`.

        Parameters
        ----------
        timeout : float
            default: -1, which means wait forever.
            Maximum time in (s) to wait before raising NotDone.

        Returns
        -------
        None
            if done before `timeout`

        Raises
        ------
        NotDone
            if `timeout` reached before I am done.
        """
        tic = time.time()
        remaining: float
        if timeout is False or timeout < 0:
            remaining = 3600 * 24 * 7  # a week
        else:
            remaining = timeout
        for evt in self.events:
            if remaining < 0:
                raise NotDone
            evt.wait(timeout=remaining)
            if not evt.is_set():
                raise NotDone
            toc = time.time()
            remaining -= toc - tic
            tic = toc

        for peer in self.peers:
            if remaining < 0:
                raise NotDone
            peer.wait(timeout=remaining)
            toc = time.time()
            remaining -= toc - tic
            tic = toc


_FINISHED_TRACKER = MessageTracker()

__all__ = ['MessageTracker', '_FINISHED_TRACKER']