From 3f268ee7c36ee0d8e4d2819679edc7c8abe3c629 Mon Sep 17 00:00:00 2001 From: Trevor Manz Date: Tue, 21 Jul 2026 11:40:09 -0400 Subject: [PATCH] Install anywidget comm provider eagerly Formatter hooks skip local modules to avoid mistaking user code for third-party packages. Editable installs look local too, so descriptor widgets could open a dummy comm before their formatter ran. These changes install the lightweight provider during formatter initialization while anywidget itself remains lazy. Refs MO-6824 --- .../formatters/anywidget_formatters.py | 7 --- marimo/_output/formatters/formatters.py | 8 +++ .../ui/_impl/anywidget/comm_provider.py | 17 ++++- .../ui/_impl/anywidget/test_comm_provider.py | 62 +++++++++++++------ 4 files changed, 64 insertions(+), 30 deletions(-) diff --git a/marimo/_output/formatters/anywidget_formatters.py b/marimo/_output/formatters/anywidget_formatters.py index 90f5e6a9dd8..2ef3752336c 100644 --- a/marimo/_output/formatters/anywidget_formatters.py +++ b/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]: diff --git a/marimo/_output/formatters/formatters.py b/marimo/_output/formatters/formatters.py index ff4222dce82..1163a396283 100644 --- a/marimo/_output/formatters/formatters.py +++ b/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. diff --git a/marimo/_plugins/ui/_impl/anywidget/comm_provider.py b/marimo/_plugins/ui/_impl/anywidget/comm_provider.py index 3c99c0a8955..9d55912857d 100644 --- a/marimo/_plugins/ui/_impl/anywidget/comm_provider.py +++ b/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,8 +42,8 @@ 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 @@ -46,12 +51,17 @@ def patch_comm_create() -> None: 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] diff --git a/tests/_plugins/ui/_impl/anywidget/test_comm_provider.py b/tests/_plugins/ui/_impl/anywidget/test_comm_provider.py index cd0484080f8..50a4295c093 100644 --- a/tests/_plugins/ui/_impl/anywidget/test_comm_provider.py +++ b/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'; } }"