Install anywidget comm provider eagerly

marimo/_output/formatters/anywidget_formatters.py

@@ -15,13 +15,6 @@ def register(self) -> None:
         import anywidget  # type: ignore [import-not-found]
 
         from marimo._output import formatting
-        from marimo._plugins.ui._impl.anywidget.comm_provider import (
-            patch_comm_create,
-        )
-
-        # Patch the comm library so anywidget's descriptor API
-        # (MimeBundleDescriptor) creates comms that work in marimo.
-        patch_comm_create()
 
         @formatting.formatter(anywidget.AnyWidget)
         def _from(lmap: anywidget.AnyWidget) -> tuple[KnownMimeType, str]:

marimo/_output/formatters/formatters.py

@@ -195,6 +195,14 @@ def register_formatters(theme: Theme = "light") -> None:
     case, the trade-off is worth it.
     """
 
+    # Descriptor-based anywidgets open their comm before formatting. Install
+    # the comm provider eagerly because local modules skip formatter hooks.
+    from marimo._plugins.ui._impl.anywidget.comm_provider import (
+        install_anywidget_comm_provider,
+    )
+
+    install_anywidget_comm_provider()
+
     # For modules that are already imported, register their formatters
     # immediately; their import hook wouldn't be triggered since they are
     # already imported. This is relevant when executing as a script.

marimo/_plugins/ui/_impl/anywidget/comm_provider.py

@@ -8,7 +8,7 @@
 
 from __future__ import annotations
 
-from typing import Any
+from typing import TYPE_CHECKING, Any
 from uuid import uuid4
 
 from marimo._loggers import marimo_logger
@@ -22,6 +22,11 @@
 
 LOGGER = marimo_logger()
 
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+_PATCHED_CREATE_COMM: Callable[..., Any] | None = None
+
 
 def _is_anywidget_comm(target_name: str, data: dict[str, Any] | None) -> bool:
     """Check if this comm is being opened by anywidget.
@@ -37,21 +42,26 @@ def _is_anywidget_comm(target_name: str, data: dict[str, Any] | None) -> bool:
     return bool(state.get("_model_module") == "anywidget")
 
 
-def patch_comm_create() -> None:
-    """Replace `comm.create_comm` with a marimo-backed implementation.
+def install_anywidget_comm_provider() -> None:
+    """Install a marimo-backed comm provider for descriptor anywidgets.
 
     Only intercepts comms created by anywidget (identified by
     `target_name` and `_model_module`). All other comms fall
     through to a no-op `DummyComm`.
 
     This is idempotent -- calling it multiple times is safe.
     """
+    global _PATCHED_CREATE_COMM
+
     try:
         import comm
         from comm import DummyComm
     except ImportError:
         return
 
+    if comm.create_comm is _PATCHED_CREATE_COMM:
+        return
+
     def _marimo_create_comm(
         *,
         target_name: str = "comm",
@@ -101,4 +111,5 @@ def _marimo_create_comm(
             pass
         return c
 
+    _PATCHED_CREATE_COMM = _marimo_create_comm
     comm.create_comm = _marimo_create_comm  # type: ignore[assignment]

tests/_plugins/ui/_impl/anywidget/test_comm_provider.py

@@ -1,6 +1,8 @@
 # Copyright 2026 Marimo. All rights reserved.
 from __future__ import annotations
 
+import sys
+
 import pytest
 
 from marimo._dependencies.dependencies import DependencyManager
@@ -15,26 +17,28 @@ def test_patches_comm_create() -> None:
         import comm
 
         from marimo._plugins.ui._impl.anywidget.comm_provider import (
-            patch_comm_create,
+            install_anywidget_comm_provider,
         )
 
         original = comm.create_comm
-        patch_comm_create()
-        assert comm.create_comm is not original
-        # Restore
-        comm.create_comm = original
+        comm.create_comm = comm.DummyComm
+        try:
+            install_anywidget_comm_provider()
+            assert comm.create_comm is not comm.DummyComm
+        finally:
+            comm.create_comm = original
 
     @staticmethod
     def test_anywidget_comm_returns_marimo_comm() -> None:
         import comm
 
         from marimo._plugins.ui._impl.anywidget.comm_provider import (
-            patch_comm_create,
+            install_anywidget_comm_provider,
         )
         from marimo._plugins.ui._impl.comm import MarimoComm
 
         original = comm.create_comm
-        patch_comm_create()
+        install_anywidget_comm_provider()
         try:
             c = comm.create_comm(
                 target_name="jupyter.widget",
@@ -64,11 +68,11 @@ def test_non_anywidget_comm_returns_dummy() -> None:
         from comm import DummyComm
 
         from marimo._plugins.ui._impl.anywidget.comm_provider import (
-            patch_comm_create,
+            install_anywidget_comm_provider,
         )
 
         original = comm.create_comm
-        patch_comm_create()
+        install_anywidget_comm_provider()
         try:
             c = comm.create_comm(
                 target_name="some.other.target",
@@ -83,18 +87,36 @@ def test_idempotent() -> None:
         import comm
 
         from marimo._plugins.ui._impl.anywidget.comm_provider import (
-            patch_comm_create,
+            install_anywidget_comm_provider,
         )
 
         original = comm.create_comm
-        patch_comm_create()
-        first = comm.create_comm
-        patch_comm_create()
-        second = comm.create_comm
-        # Second patch wraps the first — both should work
-        assert first is not original
-        assert second is not original
-        comm.create_comm = original
+        comm.create_comm = comm.DummyComm
+        try:
+            install_anywidget_comm_provider()
+            first = comm.create_comm
+            install_anywidget_comm_provider()
+            assert comm.create_comm is first
+        finally:
+            comm.create_comm = original
+
+    @staticmethod
+    def test_register_formatters_installs_provider_without_anywidget() -> None:
+        import comm
+
+        from marimo._output.formatters.formatters import register_formatters
+
+        original = comm.create_comm
+        anywidget = sys.modules.pop("anywidget", None)
+        comm.create_comm = comm.DummyComm
+        try:
+            register_formatters()
+            assert comm.create_comm is not comm.DummyComm
+            assert "anywidget" not in sys.modules
+        finally:
+            comm.create_comm = original
+            if anywidget is not None:
+                sys.modules["anywidget"] = anywidget
 
 
 @pytest.mark.skipif(not HAS_DEPS, reason="anywidget/comm not installed")
@@ -107,11 +129,11 @@ def test_produces_html_for_anywidget_mimebundle() -> None:
             _maybe_as_anywidget_html,
         )
         from marimo._plugins.ui._impl.anywidget.comm_provider import (
-            patch_comm_create,
+            install_anywidget_comm_provider,
         )
 
         original = comm.create_comm
-        patch_comm_create()
+        install_anywidget_comm_provider()
         try:
             esm = (
                 "export default { render({ el }) { el.textContent = 'hi'; } }"