improvement: use module mapping when running auto-install in micropip

marimo/_runtime/runtime.py

@@ -531,10 +531,13 @@ def _execute_stale_cells_callback(self) -> None:
         return self.enqueue_control_request(ExecuteStaleRequest())
 
     def _execute_install_missing_packages_callback(
-        self, package_manager: str
+        self, package_manager: str, packages: list[str]
     ) -> None:
+        version = {pkg: "" for pkg in packages}
         return self.enqueue_control_request(
-            InstallMissingPackagesRequest(manager=package_manager, versions={})
+            InstallMissingPackagesRequest(
+                manager=package_manager, versions=version
+            )
         )
 
     def _update_runtime_from_user_config(self, config: MarimoConfig) -> None:
@@ -902,26 +905,20 @@ def _invalidate_cell_state(
             and missing_modules_after_deletion
             != missing_modules_before_deletion
         ):
-            if self.package_manager.should_auto_install():
-                self._execute_install_missing_packages_callback(
-                    self.package_manager.name
+            packages = list(
+                sorted(
+                    pkg
+                    for mod in missing_modules_after_deletion
+                    if not self.package_manager.attempted_to_install(
+                        pkg := self.package_manager.module_to_package(mod)
+                    )
                 )
-            else:
-                # Deleting a cell can make the set of missing packages smaller
-                MissingPackageAlert(
-                    packages=list(
-                        sorted(
-                            pkg
-                            for mod in missing_modules_after_deletion
-                            if not self.package_manager.attempted_to_install(
-                                pkg := self.package_manager.module_to_package(
-                                    mod
-                                )
-                            )
-                        )
-                    ),
-                    isolated=is_python_isolated(),
-                ).broadcast()
+            )
+            # Deleting a cell can make the set of missing packages smaller
+            MissingPackageAlert(
+                packages=packages,
+                isolated=is_python_isolated(),
+            ).broadcast()
 
         cell.set_output(None)
         get_context().cell_lifecycle_registry.dispose(
@@ -1252,13 +1249,14 @@ def _broadcast_missing_packages(self, runner: cell_runner.Runner) -> None:
             ]
 
             if missing_packages:
+                packages = list(sorted(missing_packages))
                 if self.package_manager.should_auto_install():
                     self._execute_install_missing_packages_callback(
-                        self.package_manager.name
+                        self.package_manager.name, packages
                     )
                 else:
                     MissingPackageAlert(
-                        packages=list(sorted(missing_packages)),
+                        packages=packages,
                         isolated=is_python_isolated(),
                     ).broadcast()
 

tests/_runtime/test_manage_script_metadata.py

@@ -1,18 +1,29 @@
 # Copyright 2024 Marimo. All rights reserved.
 from __future__ import annotations
 
+import pathlib
+import sys
 from types import ModuleType
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, cast
 from unittest.mock import AsyncMock, call, patch
 
 import pytest
 
-from marimo._config.config import (
-    merge_default_config,
-)
+from marimo._config.config import merge_default_config
 from marimo._config.settings import GLOBAL_SETTINGS
 from marimo._dependencies.dependencies import DependencyManager
-from marimo._runtime.requests import InstallMissingPackagesRequest
+from marimo._messaging.ops import InstallingPackageAlert, MissingPackageAlert
+from marimo._runtime.packages.package_managers import create_package_manager
+from marimo._runtime.packages.pypi_package_manager import (
+    MicropipPackageManager,
+    PipPackageManager,
+)
+from marimo._runtime.packages.utils import is_python_isolated
+from marimo._runtime.requests import (
+    ControlRequest,
+    InstallMissingPackagesRequest,
+)
+from marimo._runtime.runner import cell_runner
 from tests.conftest import MockedKernel
 
 if TYPE_CHECKING:
@@ -218,8 +229,6 @@ async def test_install_missing_packages_micropip(
 ) -> None:
     k = mocked_kernel.k
     # Fake put pyodide in sys.modules
-    import sys
-
     sys.modules["pyodide"] = ModuleType("pyodide")
 
     with patch("micropip.install", new_callable=AsyncMock) as mock_install:
@@ -244,8 +253,6 @@ async def test_install_missing_packages_micropip_with_versions(
 ) -> None:
     k = mocked_kernel.k
     # Fake put pyodide in sys.modules
-    import sys
-
     sys.modules["pyodide"] = ModuleType("pyodide")
 
     with patch("micropip.install", new_callable=AsyncMock) as mock_install:
@@ -273,8 +280,6 @@ async def test_install_missing_packages_micropip_other_modules(
     k.module_registry.modules = lambda: set(
         {"idk", "done", "already_installed"}
     )
-    import sys
-
     sys.modules["pyodide"] = ModuleType("pyodide")
     sys.modules["already_installed"] = ModuleType("already_installed")
 
@@ -294,3 +299,198 @@ async def test_install_missing_packages_micropip_other_modules(
     # Remove pyodide from sys.modules
     del sys.modules["pyodide"]
     del sys.modules["already_installed"]
+
+
+async def test_broadcast_missing_packages(
+    mocked_kernel: MockedKernel,
+) -> None:
+    """Test that _broadcast_missing_packages correctly handles missing packages for micropip"""
+    k = mocked_kernel.k
+    control_requests: list[ControlRequest] = []
+    broadcast_messages: list[InstallingPackageAlert | MissingPackageAlert] = []
+
+    def mock_enqueue(request: ControlRequest) -> None:
+        control_requests.append(request)
+
+    def mock_broadcast(
+        msg: InstallingPackageAlert | MissingPackageAlert,
+    ) -> None:
+        broadcast_messages.append(msg)
+
+    k.enqueue_control_request = mock_enqueue
+    InstallingPackageAlert.broadcast = mock_broadcast  # type: ignore
+    MissingPackageAlert.broadcast = mock_broadcast  # type: ignore
+
+    sys.modules["pyodide"] = ModuleType("pyodide")
+
+    # Create a mock runner with ModuleNotFoundError
+    class MockRunner:
+        def __init__(self) -> None:
+            self.exceptions = {
+                "cell1": ModuleNotFoundError(
+                    "No module named 'numpy'", name="numpy"
+                ),
+                # Duplicate
+                "cell2": ModuleNotFoundError(
+                    "No module named 'numpy'", name="numpy"
+                ),
+                # Has mapping
+                "ibis": ModuleNotFoundError(
+                    "No module named 'ibis'", name="ibis"
+                ),
+            }
+
+    # Case 1: Auto-install enabled
+    with patch("micropip.install", new_callable=AsyncMock):
+        runner = cast(cell_runner.Runner, MockRunner())
+        k.package_manager = create_package_manager("micropip")
+        package_manager = k.package_manager
+        assert isinstance(package_manager, MicropipPackageManager)
+        k._broadcast_missing_packages(runner)
+
+        # Should create install request
+        assert len(control_requests) == 1
+        request = control_requests[0]
+        assert isinstance(request, InstallMissingPackagesRequest)
+        assert request.manager == package_manager.name
+        assert request.versions == {"numpy": "", "ibis-framework[duckdb]": ""}
+        assert len(broadcast_messages) == 0
+
+        # Case 2: Auto-install disabled
+        control_requests.clear()
+        broadcast_messages.clear()
+        package_manager.should_auto_install = lambda: False  # type: ignore
+        k._broadcast_missing_packages(runner)
+
+        # Should broadcast alert instead of installing
+        assert len(control_requests) == 0
+        assert len(broadcast_messages) == 1
+        alert = broadcast_messages[0]
+        assert isinstance(alert, MissingPackageAlert)
+        assert alert.packages == ["ibis-framework[duckdb]", "numpy"]
+        assert alert.isolated == is_python_isolated()
+
+        # Case 3: Multiple missing modules
+        control_requests.clear()
+        broadcast_messages.clear()
+        k.module_registry.missing_modules = lambda: {
+            "ibis-framework[duckdb]",
+            "pandas",
+            "scipy",
+        }  # type: ignore
+        package_manager.should_auto_install = lambda: True  # type: ignore
+        k._broadcast_missing_packages(runner)
+
+        # Should create install request with all missing packages
+        assert len(control_requests) == 1
+        request = control_requests[0]
+        assert isinstance(request, InstallMissingPackagesRequest)
+        assert request.manager == package_manager.name
+        assert request.versions == {
+            "ibis-framework[duckdb]": "",
+            "numpy": "",
+            "pandas": "",
+            "scipy": "",
+        }
+
+        # Case 4: Already attempted packages should be filtered
+        control_requests.clear()
+        broadcast_messages.clear()
+        package_manager.attempted_to_install = (
+            lambda package: package == "numpy"
+        )  # type: ignore
+        k._broadcast_missing_packages(runner)
+
+        # Should only include packages not yet attempted
+        assert len(control_requests) == 1
+        request = control_requests[0]
+        assert isinstance(request, InstallMissingPackagesRequest)
+        assert request.manager == package_manager.name
+        assert request.versions == {
+            "ibis-framework[duckdb]": "",
+            "pandas": "",
+            "scipy": "",
+        }
+
+    del sys.modules["pyodide"]
+
+
+def test_broadcast_missing_packages_pip(
+    mocked_kernel: MockedKernel,
+) -> None:
+    """Test that _broadcast_missing_packages correctly handles missing packages for pip"""
+    k = mocked_kernel.k
+    control_requests: list[ControlRequest] = []
+    broadcast_messages: list[InstallingPackageAlert | MissingPackageAlert] = []
+
+    def mock_enqueue(request: ControlRequest) -> None:
+        control_requests.append(request)
+
+    def mock_broadcast(
+        msg: InstallingPackageAlert | MissingPackageAlert,
+    ) -> None:
+        broadcast_messages.append(msg)
+
+    k.enqueue_control_request = mock_enqueue
+    InstallingPackageAlert.broadcast = mock_broadcast  # type: ignore
+    MissingPackageAlert.broadcast = mock_broadcast  # type: ignore
+
+    # Create a mock runner with ModuleNotFoundError
+    class MockRunner:
+        def __init__(self) -> None:
+            self.exceptions = {
+                "cell1": ModuleNotFoundError(
+                    "No module named 'numpy'", name="numpy"
+                ),
+                # Duplicate
+                "cell2": ModuleNotFoundError(
+                    "No module named 'numpy'", name="numpy"
+                ),
+                # Has mapping
+                "ibis": ModuleNotFoundError(
+                    "No module named 'ibis'", name="ibis"
+                ),
+            }
+
+    k.package_manager = create_package_manager("pip")
+    package_manager = k.package_manager
+    assert isinstance(package_manager, PipPackageManager)
+    package_manager.install = AsyncMock()
+    runner = cast(cell_runner.Runner, MockRunner())
+
+    # Case 1: Missing modules with auto-install disabled
+    k.module_registry.missing_modules = lambda: {"numpy", "pandas"}  # type: ignore
+    package_manager.should_auto_install = lambda: False  # type: ignore
+    k._broadcast_missing_packages(runner)
+
+    # Should broadcast alert instead of installing
+    assert len(control_requests) == 0
+    assert len(broadcast_messages) == 1
+    alert = broadcast_messages[0]
+    assert isinstance(alert, MissingPackageAlert)
+    assert alert.packages == ["ibis-framework[duckdb]", "numpy", "pandas"]
+    assert alert.isolated == is_python_isolated()
+
+    # Case 2: Multiple missing modules with auto-install enabled
+    control_requests.clear()
+    broadcast_messages.clear()
+    k.module_registry.missing_modules = lambda: {
+        "ibis-framework[duckdb]",
+        "numpy",
+        "pandas",
+        "scipy",
+    }  # type: ignore
+    package_manager.should_auto_install = lambda: True  # type: ignore
+    k._broadcast_missing_packages(runner)
+
+    # Should create install request with all missing packages
+    assert len(control_requests) == 1
+    request = control_requests[0]
+    assert isinstance(request, InstallMissingPackagesRequest)
+    assert request.manager == "pip"
+    assert request.versions == {
+        "ibis-framework[duckdb]": "",
+        "numpy": "",
+        "pandas": "",
+        "scipy": "",
+    }