"""Client for Qt's WebChannel protocol over a plain websocket, without any Qt dependency.
Some instruments (for example Menlo Systems frequency combs) expose their control/status tree via Qt's
`QWebChannel <https://doc.qt.io/qt-5/qtwebchannel-index.html>`_ protocol. The protocol implementation itself
is provided by the :mod:`pywebchannel` package, which has no Qt dependency. This module only supplies the
websocket transport and the endpoint's authentication handshake on top of it.
"""
import json
import threading
import time
import websocket
from heros.helper import log
try:
from pywebchannel.qwebchannel import QWebChannel
except ModuleNotFoundError:
QWebChannel = None # type: ignore[assignment,misc]
log.exception(
"Could not import the 'pywebchannel' module, required for QWebChannel connections. "
"Install it with `uv pip install git+https://github.com/MenloSystems/pywebchannel`"
)
QWEBCHANNEL_DEFAULT_PORT = 8002
# websocket-client has no per-connection connect timeout; the only way to bound a socket connect is the
# process-wide websocket.setdefaulttimeout(). This lock serializes the set-then-connect window across every
# QWebChannelConnection instance so two connections with different `timeout` values can't race on it.
_connect_timeout_lock = threading.Lock()
[docs]
class QWebChannelConnection:
"""Manage a QWebChannel connection to a device over a plain websocket.
Args:
host: Hostname or IP address of the websocket endpoint.
port: Port the websocket endpoint listens on.
path: Path component of the websocket URL.
user: Username used for authentication.
password: Password used for authentication.
timeout: Seconds to wait for the connection to open and the channel to initialize. Also bounds how
long the underlying socket connect may block, so an unresponsive (as opposed to actively
refusing) endpoint fails within `timeout` too, instead of hanging on the OS's own connect timeout.
reconnect_cooldown: Minimum seconds between two connection attempts. A :py:meth:`connect` call made
before this has elapsed since the last attempt fails immediately without touching the network,
so a caller that retries on every use (like :class:`OFC`) does not hammer a device that
stays unreachable for an extended period.
Note:
There is no active reconnect-on-disconnect: a lost connection is only noticed and re-established the
next time :py:meth:`connect` is called (which every read/write on :class:`OFC` does implicitly
via its `_ensure_connected` helper). This keeps reconnection logic in one place, serialized by a lock,
instead of racing an active retry from the websocket's own background thread against callers.
"""
def __init__(
self,
host: str,
port: int = QWEBCHANNEL_DEFAULT_PORT,
path: str = "/core/",
user: str = "guest",
password: str = "",
timeout: float = 5.0,
reconnect_cooldown: float = 30.0,
) -> None:
self.host = host
self.port = port
self.path = path
self.user = user
self.password = password
self.timeout = timeout
self.reconnect_cooldown = reconnect_cooldown
self.channel = QWebChannel(initCallback=self._on_ready) # type: ignore
self._ws: websocket.WebSocketApp | None = None
self._thread: threading.Thread | None = None
self._authenticated = False
self._ready = threading.Event()
self._connect_lock = threading.Lock()
self._last_attempt: float | None = None
@property
def root(self): # noqa: ANN201
"""The root object exposed by the QWebChannel endpoint, or None if not connected yet."""
return self.channel.objects.get("root")
[docs]
def is_ready(self) -> bool:
"""Whether the channel is initialized and :py:attr:`root` is available."""
return self._ready.is_set()
[docs]
def send(self, data: str) -> None:
"""Send raw data over the websocket.
Called by :py:class:`pywebchannel.qwebchannel.QWebChannel`, not meant to be called directly.
"""
if self._ws is None:
msg = f"Not connected to {self.host}:{self.port}"
raise ConnectionError(msg)
self._ws.send(data)
[docs]
def connect(self) -> None:
"""Open the websocket connection and wait until the channel is initialized.
Does nothing if a connection is already open and ready. Safe to call concurrently from multiple
threads: only the first caller actually (re)connects, the others wait for it and then observe the
result.
Raises:
ConnectionError: The last attempt was less than `reconnect_cooldown` seconds ago; no new attempt
was made.
TimeoutError: The connection could not be established within `timeout` seconds.
"""
with self._connect_lock:
if self.is_ready():
return
now = time.monotonic()
if self._last_attempt is not None and now - self._last_attempt < self.reconnect_cooldown:
msg = (
f"Not retrying connection to {self.host}:{self.port} yet, "
f"last attempt was {now - self._last_attempt:.0f}s ago (cooldown {self.reconnect_cooldown:.0f}s)"
)
raise ConnectionError(msg)
self._last_attempt = now
if self._ws is not None:
# a previous attempt left a socket/thread behind (e.g. it timed out but kept running); close
# it so it doesn't leak and so its callbacks stop touching this connection's state
self._ws.close()
self._ready.clear()
self._authenticated = False
url = f"ws://{self.host}:{self.port}{self.path}"
self._ws = websocket.WebSocketApp(
url,
on_open=self._on_open,
on_message=self._on_message,
on_error=self._on_error,
on_close=self._on_close,
)
with _connect_timeout_lock:
# bounds the underlying socket connect (and any later blocking send/recv) to `timeout`, so an
# unresponsive endpoint fails within `timeout` too, not just an actively refusing one. This is
# a process-wide default (websocket-client has no per-connection timeout), so the lock is held
# until this connection is either ready or has given up, keeping another connection's
# concurrent connect() from silently overwriting the timeout applied to this one
websocket.setdefaulttimeout(self.timeout)
self._thread = threading.Thread(target=self._ws.run_forever, daemon=True)
self._thread.start()
if not self._ready.wait(self.timeout):
msg = f"Timed out connecting to QWebChannel endpoint {url}"
raise TimeoutError(msg)
[docs]
def close(self) -> None:
"""Close the websocket connection."""
if self._ws is not None:
self._ws.close()
self._ready.clear()
def _on_open(self, ws: websocket.WebSocketApp) -> None:
log.debug("Connected to QWebChannel endpoint %s:%s, authenticating", self.host, self.port)
ws.send(json.dumps({"user": self.user, "password": self.password}))
def _on_message(self, ws: websocket.WebSocketApp, message: str) -> None:
if ws is not self._ws:
# a stale callback from a connection attempt that has since been superseded/closed
return
if not self._authenticated:
try:
data = json.loads(message)
except json.JSONDecodeError:
log.warning("Received non-json message before authentication: %s", message)
return
if data.get("authenticated") is True:
self._authenticated = True
self.channel.connection_made(self)
else:
log.warning("Authentication failed for QWebChannel endpoint %s:%s", self.host, self.port)
return
self.channel.message_received(message)
def _on_error(self, _ws: websocket.WebSocketApp, error: Exception) -> None:
log.error("QWebChannel connection error on %s:%s: %s", self.host, self.port, error)
def _on_close(self, ws: websocket.WebSocketApp, *_args) -> None:
if ws is not self._ws:
# a stale callback from a connection attempt that has since been superseded/closed
return
log.debug("QWebChannel connection to %s:%s closed", self.host, self.port)
self._ready.clear()
self.channel.connection_closed()
def _on_ready(self, _channel: object) -> None:
self._ready.set()