Improve type hints

examples/widgets.ipynb

@@ -39,7 +39,9 @@
     "\n",
     "import ipywidgets as ipw\n",
     "\n",
-    "import ipylab"
+    "import ipylab\n",
+    "\n",
+    "app = ipylab.JupyterFrontEnd()"
    ]
   },
   {
@@ -743,24 +745,6 @@
     "t = n.to_task(update())"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "t = app.notification.new_action(\"Some action\", lambda: app.commands.execute(\"help:about\"))"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "action = t.result()"
-   ]
-  },
   {
    "cell_type": "markdown",
    "metadata": {},

ipylab/asyncwidget.py

@@ -6,10 +6,10 @@
 import inspect
 import traceback
 import uuid
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, TypeVar
 
 from ipywidgets import Widget, register, widget_serialization
-from traitlets import Container, Dict, Instance, Set, Unicode
+from traitlets import Container, Dict, HasTraits, Instance, Set, Unicode
 
 import ipylab._frontend as _fe
 from ipylab.common import JavascriptType, Transform, TransformType
@@ -25,6 +25,8 @@
 
 __all__ = ["AsyncWidgetBase", "WidgetBase", "register", "pack", "Widget"]
 
+T = TypeVar("T")
+
 
 def pack(obj: Widget | Any):
     """Return serialized obj if it is a Widget otherwise return it unchanged."""
@@ -128,8 +130,7 @@ def _check_closed(self):
             msg = f"This widget is closed {self!r}"
             raise RuntimeError(msg)
 
-    # TODO: Add better type hints (pass the result of the coro)
-    def to_task(self, coro: Coroutine):
+    def to_task(self, coro: Coroutine[None, None, T]) -> Task[T]:
         """Run the coro in a task."""
 
         self._check_closed()
@@ -138,8 +139,7 @@ def to_task(self, coro: Coroutine):
         task.add_done_callback(self._tasks.discard)
         return task
 
-    # TODO: Add better type hints (pass the result of the coro)
-    async def _wrap_coro(self, coro: Coroutine):
+    async def _wrap_coro(self, coro: Coroutine[None, None, T]) -> T:
         try:
             return await coro
         except asyncio.CancelledError:
@@ -172,14 +172,14 @@ def _check_get_error(self, content: dict | None = None) -> IpylabFrontendError |
             return IpylabFrontendError(f'{self.__class__.__name__} failed with message "{error}"')
         return None
 
-    async def _add_to_tuple_trait(self, name: str, item: Awaitable[Widget] | Widget):
+    async def _add_to_tuple_trait(self, name: str, item: Awaitable[T] | T) -> T:
         """Add the item to the tuple and observe its comm."""
         if inspect.isawaitable(item):
-            value = await item
+            value: T = await item
         else:
-            value = item
+            value: T = item
         items = getattr(self, name)
-        if value not in items:
+        if isinstance(value, HasTraits) and value not in items:
             value.observe(lambda _: self.set_trait(name, tuple(i for i in getattr(self, name) if i.comm)), "comm")
             self.set_trait(name, (*items, value))
         return value

ipylab/connection.py

@@ -12,7 +12,6 @@
 from ipylab.asyncwidget import AsyncWidgetBase
 
 if TYPE_CHECKING:
-    from asyncio import Task
     from collections.abc import Generator
     from typing import Literal, overload
 
@@ -146,7 +145,7 @@ def get_existing_connection(cls, *name_or_id: str, quiet=False):
 class MainAreaConnection(Connection):
     CID_PREFIX = "ipylab MainArea"
 
-    def activate(self) -> Task[Self]:
+    def activate(self):
         self._check_closed()
 
         async def activate_():

ipylab/notification.py

@@ -88,8 +88,8 @@ async def _do_operation_for_frontend(self, operation: str, payload: dict, buffer
         """Overload this function as required."""
         match operation:
             case "action callback":
-                ActionConnection.get_existing_connection(payload["cid"]).callback()
-        await super()._do_operation_for_frontend(operation, payload, buffers)
+                return ActionConnection.get_existing_connection(payload["cid"]).callback()
+        return await super()._do_operation_for_frontend(operation, payload, buffers)
 
     async def _ensure_action(self, value: ActionConnection | NotifyAction) -> ActionConnection:
         "Create a new action."
@@ -130,9 +130,10 @@ async def notify():
             )
             for action in actions:
                 await notification._add_to_tuple_trait("actions", action)  # noqa: SLF001
+            await self._add_to_tuple_trait("notifications", notification)
             return notification
 
-        return self.to_task(self._add_to_tuple_trait("notifications", notify()))
+        return self.to_task(notify())
 
     def new_action(
         self,