fix: resolve normalized package names to their module

marimo/_runtime/packages/package_manager.py

@@ -2,6 +2,7 @@
 from __future__ import annotations
 
 import abc
+import re
 import subprocess
 import sys
 from collections.abc import Callable
@@ -28,6 +29,16 @@
 LogCallback = Callable[[str], None]
 
 
+def _normalize_package_name(name: str) -> str:
+    """Normalize a package name per PEP 503.
+
+    PyPI treats package names as case-insensitive and considers runs of
+    `-`, `_`, and `.` as equivalent, so `Pillow`, `pillow`, and `scikit_learn`
+    all normalize to a canonical form.
+    """
+    return re.sub(r"[-_.]+", "-", name).lower()
+
+
 class PackageDescription(msgspec.Struct, rename="camel"):
     name: str
     version: str
@@ -254,6 +265,10 @@ def __init__(self, python_exe: str | None = None) -> None:
         # Initialized lazily
         self._module_name_to_repo_name: dict[str, str] | None = None
         self._repo_name_to_module_name: dict[str, str] | None = None
+        # Reverse map keyed on PEP 503-normalized package names, so we can
+        # resolve names that uv has normalized (e.g. `Pillow` -> `pillow`)
+        # back to their module.
+        self._normalized_repo_name_to_module_name: dict[str, str] | None = None
         # Python executable for targeting a specific venv (used by pip/uv)
         # Defaults to sys.executable if not provided
         self._python_exe = python_exe or PY_EXE
@@ -273,6 +288,12 @@ def _initialize_mappings(self) -> None:
                 v: k for k, v in self._module_name_to_repo_name.items()
             }
 
+        if self._normalized_repo_name_to_module_name is None:
+            self._normalized_repo_name_to_module_name = {
+                _normalize_package_name(k): v
+                for k, v in self._repo_name_to_module_name.items()
+            }
+
     def module_to_package(self, module_name: str) -> str:
         """Canonicalizes a module name to a package name on PyPI."""
         if self._module_name_to_repo_name is None:
@@ -286,12 +307,24 @@ def module_to_package(self, module_name: str) -> str:
 
     def package_to_module(self, package_name: str) -> str:
         """Canonicalizes a package name to a module name."""
-        if self._repo_name_to_module_name is None:
+        if (
+            self._repo_name_to_module_name is None
+            or self._normalized_repo_name_to_module_name is None
+        ):
             self._initialize_mappings()
         assert self._repo_name_to_module_name is not None
+        assert self._normalized_repo_name_to_module_name is not None
 
-        return (
-            self._repo_name_to_module_name[package_name]
-            if package_name in self._repo_name_to_module_name
-            else package_name.replace("-", "_")
-        )
+        # Exact match first, to preserve any casing in the known mapping.
+        if package_name in self._repo_name_to_module_name:
+            return self._repo_name_to_module_name[package_name]
+
+        # PyPI package names are case-insensitive and treat runs of `-`, `_`,
+        # `.` as equivalent (PEP 503). uv normalizes names when it writes them
+        # into a notebook's script metadata (e.g. `Pillow` -> `pillow`), so
+        # fall back to a normalized lookup before guessing.
+        normalized = _normalize_package_name(package_name)
+        if normalized in self._normalized_repo_name_to_module_name:
+            return self._normalized_repo_name_to_module_name[normalized]
+
+        return package_name.replace("-", "_")

tests/_runtime/packages/test_pypi_package_manager.py

@@ -57,6 +57,20 @@ def test_package_to_module() -> None:
     assert mgr.package_to_module("scikit-learn") == "sklearn"
 
 
+def test_package_to_module_is_normalized() -> None:
+    # PyPI package names are case-insensitive and treat runs of `-`, `_`, `.`
+    # as equivalent (PEP 503). uv normalizes names when it writes them into a
+    # notebook's script metadata (e.g. `Pillow` -> `pillow`), so the reverse
+    # mapping must still resolve them back to the correct module.
+    # Regression test for https://github.com/marimo-team/marimo/issues/9801
+    mgr = PipPackageManager()
+    assert mgr.package_to_module("Pillow") == "PIL"
+    assert mgr.package_to_module("pillow") == "PIL"
+    assert mgr.package_to_module("scikit-learn") == "sklearn"
+    assert mgr.package_to_module("Scikit-Learn") == "sklearn"
+    assert mgr.package_to_module("scikit_learn") == "sklearn"
+
+
 async def test_failed_install_returns_false() -> None:
     mgr = PipPackageManager()
     # almost surely does not exist