diff --git a/marimo/_runtime/packages/package_manager.py b/marimo/_runtime/packages/package_manager.py index b38178468d5..8de7d336f08 100644 --- a/marimo/_runtime/packages/package_manager.py +++ b/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("-", "_") diff --git a/tests/_runtime/packages/test_pypi_package_manager.py b/tests/_runtime/packages/test_pypi_package_manager.py index b388c0879bc..a9673c25576 100644 --- a/tests/_runtime/packages/test_pypi_package_manager.py +++ b/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