Support awareness over Comm

Author: gjmooney

src/ypywidgets/comm.py

@@ -1,7 +1,9 @@
 from __future__ import annotations
+from typing import Any, Callable
 
 import comm
 from pycrdt import (
+    Awareness,
     Doc,
     Text,
     TransactionEvent,
@@ -10,6 +12,7 @@
     create_sync_message,
     create_update_message,
     handle_sync_message,
+    read_message,
 )
 
 from .widget import Widget
@@ -48,18 +51,27 @@ def __init__(
     ) -> None:
         self._ydoc = ydoc
         self._comm = comm
+        self._awareness = Awareness(ydoc)
         msg = create_sync_message(ydoc)
         self._comm.send(buffers=[msg])
         self._comm.on_msg(self._receive)
 
+    @property
+    def awareness(self) -> Awareness:
+        return self._awareness
+
     def _receive(self, msg):
         message = bytes(msg["buffers"][0])
-        if message[0] == YMessageType.SYNC:
+        mtype = message[0]
+        if mtype == YMessageType.SYNC:
             reply = handle_sync_message(message[1:], self._ydoc)
             if reply is not None:
                 self._comm.send(buffers=[reply])
             if message[1] == YSyncMessageType.SYNC_STEP2:
                 self._ydoc.observe(self._send)
+        elif mtype == YMessageType.AWARENESS:
+            payload = read_message(message[1:])
+            self._awareness.apply_awareness_update(payload, None)
 
     def _send(self, event: TransactionEvent):
         update = event.update
@@ -69,12 +81,12 @@ def _send(self, event: TransactionEvent):
 
 class CommWidget(Widget):
     def __init__(
-            self,
-            ydoc: Doc | None = None,
-            comm_data: dict | None = None,
-            comm_metadata: dict | None = None,
-            comm_id: str | None = None,
-        ):
+        self,
+        ydoc: Doc | None = None,
+        comm_data: dict | None = None,
+        comm_metadata: dict | None = None,
+        comm_id: str | None = None,
+    ):
         super().__init__(ydoc)
         model_name = self.__class__.__name__
         _model_name = self.ydoc["_model_name"] = Text()
@@ -85,7 +97,21 @@ def __init__(
                 create_ydoc=not ydoc,
             )
         self._comm = create_widget_comm(comm_data, comm_metadata, comm_id)
-        CommProvider(self.ydoc, self._comm)
+        self._comm_provider = CommProvider(self.ydoc, self._comm)
+
+    @property
+    def awareness(self) -> Awareness:
+        return self._comm_provider.awareness
+
+    def on_awareness_change(
+        self,
+        callback: Callable[[str, tuple[dict[str, Any], Any]], None],
+    ) -> str:
+        """Subscribe to pycrdt Awareness updates; returns subscription id for unobserve."""
+        return self.awareness.observe(callback)
+
+    def unobserve_awareness(self, subscription_id: str) -> None:
+        self.awareness.unobserve(subscription_id)
 
     def _repr_mimebundle_(self, *args, **kwargs):  # pragma: nocover
         plaintext = repr(self)

tests/conftest.py

@@ -5,12 +5,14 @@
 import comm
 import pytest
 from pycrdt import (
+    Awareness,
     YMessageType,
     YSyncMessageType,
     TransactionEvent,
     create_sync_message,
     create_update_message,
     handle_sync_message,
+    read_message,
 )
 from ypywidgets import Widget
 from ypywidgets.comm import CommWidget
@@ -67,6 +69,7 @@ def __init__(self, widget_factory, comm):
         self.widget_factory = widget_factory
         self.comm = comm
         self.widget = None
+        self._remote_awareness: Optional[Awareness] = None
         self.receive_task = asyncio.create_task(self.receive())
 
     def send(self, event: TransactionEvent):
@@ -89,6 +92,11 @@ async def receive(self):
                         self.comm.handle_msg({"buffers": [reply]})
                     if message[1] == YSyncMessageType.SYNC_STEP2:
                         self.widget.ydoc.observe(self.send)
+                elif message[0] == YMessageType.AWARENESS:
+                    if self._remote_awareness is None:
+                        self._remote_awareness = Awareness(self.widget.ydoc)
+                    payload = read_message(bytes(message[1:]))
+                    self._remote_awareness.apply_awareness_update(payload, None)
 
     async def get_widget(self, timeout=0.1):
         t = time.monotonic()