Source code for ezmsg.qt.subscriber

"""EzSubscriber - Receive messages from ezmsg topics in Qt widgets."""

from __future__ import annotations

from collections.abc import Callable
from enum import Enum
from typing import Any
from typing import TYPE_CHECKING

from qtpy import QtCore
from qtpy import QtWidgets

from .sidecar import normalize_topic

if TYPE_CHECKING:
    from ezmsg.core.subclient import Subscriber

    from .session import EzSession


TopicLike = Enum | str | None




class EzSubscriber(QtCore.QObject):
    """
    Receive messages from an ezmsg topic.

    EzSubscriber provides a Qt-native interface for subscribing to ezmsg topics.
    Messages are delivered via a Qt signal, allowing seamless integration with
    Qt's event loop and signal/slot mechanism.

    For processing graphs, use ProcessorGraph directly instead.

    Example:
        class MyWidget(QtWidgets.QWidget):
            def __init__(self):
                super().__init__()
                # Simple subscription
                self.velocity_data = EzSubscriber(VelocityTopic.OUTPUT, parent=self)
                self.velocity_data.connect(self.on_velocity)

            def on_velocity(self, msg):
                self.plot.update(msg)
    """

    received = QtCore.Signal(object)  # pyright: ignore[reportPrivateImportUsage]
    switch_started = QtCore.Signal(object)  # pyright: ignore[reportPrivateImportUsage]
    topic_changed = QtCore.Signal(object)  # pyright: ignore[reportPrivateImportUsage]
    topic_cleared = QtCore.Signal()  # pyright: ignore[reportPrivateImportUsage]
    switch_failed = QtCore.Signal(object, str)  # pyright: ignore[reportPrivateImportUsage]



    def __init__(
        self,
        topic: TopicLike = None,
        parent: QtCore.QObject | None = None,
        *,
        session: EzSession | None = None,
        auto_gate: bool = False,
        leaky: bool = False,
        max_queue: int | None = None,
        throttle_hz: float | None = None,
    ):
        """
        Create a subscriber for an ezmsg topic.

        Args:
            topic: The initial topic to subscribe to, or None to start unsubscribed.
            parent: Optional parent QObject for lifecycle management.
            auto_gate: If True, suppress delivery while the parent widget is hidden.
            leaky: If True, the underlying ezmsg Subscriber will drop old messages
                if the receiver can't keep up (no backpressure).
            max_queue: Queue depth for leaky mode. If None, ezmsg defaults apply.
            throttle_hz: If set, throttle delivery by reading at most this many
                messages per second from the underlying ezmsg Subscriber.
        """
        super().__init__(parent)
        self._topic = self._validate_topic(topic)
        self._desired_topic = self._topic
        self._sub: Subscriber | None = None  # Set by EzSession during setup
        self._session: EzSession | None = None
        self._auto_gate = bool(auto_gate)
        self._leaky = bool(leaky)
        self._max_queue = max_queue
        self._throttle_hz = throttle_hz
        self._emit_epoch = 0
        self._visibility_filter: QtCore.QObject | None = None

        if (
            self._auto_gate
            and parent is not None
            and not isinstance(parent, QtWidgets.QWidget)
        ):
            raise TypeError("EzSubscriber auto_gate requires a QWidget parent")

        if self._max_queue is not None and self._max_queue <= 0:
            raise ValueError("max_queue must be positive")
        if self._throttle_hz is not None and self._throttle_hz <= 0:
            raise ValueError("throttle_hz must be positive")

        if session is not None:
            session.attach(self)


    @property
    def topic(self) -> TopicLike:
        """The currently active topic, or None if unsubscribed."""
        return self._topic

    @property
    def session(self) -> EzSession | None:
        """The session this subscriber is attached to, if any."""
        return self._session

    @property
    def auto_gate(self) -> bool:
        """Whether delivery is suppressed while the parent widget is hidden."""
        return self._auto_gate

    @property
    def leaky(self) -> bool:
        """Whether this subscriber drops old messages instead of backpressure."""
        return self._leaky

    @property
    def max_queue(self) -> int | None:
        """Leaky notification queue depth (ignored if leaky=False)."""
        return self._max_queue

    @property
    def throttle_hz(self) -> float | None:
        """If set, throttle reads from the underlying ezmsg Subscriber."""
        return self._throttle_hz

    @property
    def parent_widget(self) -> QtWidgets.QWidget | None:
        """The QWidget used for auto-gating, if available."""
        parent = self.parent()
        if isinstance(parent, QtWidgets.QWidget):
            return parent
        return None



    def connect(self, slot: Callable[[Any], None]) -> None:
        """
        Connect a handler to receive messages.

        Args:
            slot: A callable that will be invoked with each received message.
        """
        self.received.connect(slot)




    def set_topic(self, topic: Enum | str) -> None:
        """Switch to a new topic on a running session.

        This call blocks until the session applies the topic change.
        """
        validated = self._validate_topic(topic)
        session = self._require_running_session()

        previous_desired = self._desired_topic
        self._desired_topic = validated
        self.switch_started.emit(validated)
        try:
            session._set_subscriber_topic(self, validated)
            self.topic_changed.emit(self._topic)
        except Exception as exc:
            self._desired_topic = previous_desired
            self.switch_failed.emit(validated, str(exc))
            raise




    def clear_topic(self) -> None:
        """Unsubscribe from the current topic on a running session.

        This call blocks until the session disconnects the route.
        """
        session = self._require_running_session()

        previous_desired = self._desired_topic
        self._desired_topic = None
        self.switch_started.emit(None)
        try:
            session._set_subscriber_topic(self, None)
            self.topic_cleared.emit()
        except Exception as exc:
            self._desired_topic = previous_desired
            self.switch_failed.emit(None, str(exc))
            raise


    @QtCore.Slot(object, int)  # pyright: ignore[reportPrivateImportUsage]
    def _on_message(self, msg: Any, epoch: int) -> None:
        """Internal slot called from the Qt signal dispatcher."""
        if epoch != self._emit_epoch:
            return
        self.received.emit(msg)

    def _bind_session(self, session: EzSession) -> None:
        if self._session is not None and self._session is not session:
            raise RuntimeError(
                "EzSubscriber is already attached to a different session"
            )
        self._session = session

    def _require_running_session(self) -> EzSession:
        if self._session is None:
            raise RuntimeError(
                "EzSubscriber must be attached to a session before switching"
            )
        if not self._session.running:
            raise RuntimeError(
                "EzSubscriber topic switching requires a running EzSession"
            )
        return self._session

    @staticmethod
    def _validate_topic(topic: TopicLike) -> TopicLike:
        if topic is None:
            return None
        normalize_topic(topic)
        return topic