From 87700e9143af8600951b20659f01f856ca31bfa4 Mon Sep 17 00:00:00 2001 From: Myles Scolnick Date: Mon, 13 Jul 2026 14:44:23 -0400 Subject: [PATCH] fix: resolve normalized package names to their module MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Closes #9801 ## Problem `PackageManager.package_to_module` was case-sensitive. PyPI package names, however, are case-insensitive and treat runs of `-`, `_`, and `.` as equivalent (PEP 503). uv normalizes names when it writes them into a notebook's PEP 723 script metadata — e.g. `Pillow` becomes `pillow` — so the reverse lookup silently failed: ```python mgr.package_to_module("pillow") # -> "pillow" (should be "PIL") mgr.package_to_module("Scikit-Learn") # -> "Scikit_Learn" (should be "sklearn") ``` In a sandbox this is the root cause of the confusing PIL/Pillow behavior in the issue: after marimo installs `pillow`, it maps the installed package name back to a module to decide which cells to re-run. Because `pillow` resolved to the module `pillow` (which nothing imports) instead of `PIL`, the cell that imports `PIL` was never re-run, forcing the user to manually toggle the import every time they reopened the notebook. ## Fix `package_to_module` now falls back to a PEP 503-normalized lookup when there's no exact match, before guessing. The known-name mapping is still consulted exactly first, so existing casing is preserved. ```python mgr.package_to_module("pillow") # -> "PIL" mgr.package_to_module("PILLOW") # -> "PIL" mgr.package_to_module("Scikit-Learn") # -> "sklearn" mgr.package_to_module("scikit_learn") # -> "sklearn" ``` Unknown packages still fall back to the previous `-` → `_` behavior. --- marimo/_runtime/packages/package_manager.py | 45 ++++++++++++++++--- .../packages/test_pypi_package_manager.py | 14 ++++++ 2 files changed, 53 insertions(+), 6 deletions(-) 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