feat: add pypi-to-conda-name overrides to pyproject parsing

README.md

@@ -336,6 +336,18 @@ platforms = [
 ]
 ```
 
+#### PyPI name mapping
+
+If you would like to supplement or override the pypi-to-conda name mappings provided by
+[pypi-mapping][mapping], you can do so by adding a `pypi-to-conda-name` section:
+
+```toml
+# pyproject.toml
+
+[tool.conda-lock.pypi-to-conda-name]
+cupy-cuda11x = "cupy"
+```
+
 #### Extras
 
 If your pyproject.toml file contains optional dependencies/extras these can be referred to by using the `--extras` flag

conda_lock/lookup.py

@@ -1,6 +1,9 @@
+import logging
+
+from contextlib import suppress
 from functools import cached_property
 from pathlib import Path
-from typing import Dict
+from typing import Dict, Union, cast
 
 import requests
 import yaml
@@ -9,89 +12,99 @@
 from typing_extensions import TypedDict
 
 
+DEFAULT_MAPPING_URL = "https://raw.githubusercontent.com/regro/cf-graph-countyfair/master/mappings/pypi/grayskull_pypi_mapping.yaml"
+
+
 class MappingEntry(TypedDict):
     conda_name: str
-    # legacy field, generally not used by anything anymore
-    conda_forge: str
     pypi_name: NormalizedName
 
 
 class _LookupLoader:
-    _mapping_url: str = "https://raw.githubusercontent.com/regro/cf-graph-countyfair/master/mappings/pypi/grayskull_pypi_mapping.yaml"
+    """Object used to map PyPI package names to conda names."""
 
-    @property
-    def mapping_url(self) -> str:
-        return self._mapping_url
-
-    @mapping_url.setter
-    def mapping_url(self, value: str) -> None:
-        if self._mapping_url != value:
-            self._mapping_url = value
-            # Invalidate cache
-            try:
-                del self.pypi_lookup
-            except AttributeError:
-                pass
-            try:
-                del self.conda_lookup
-            except AttributeError:
-                pass
+    mapping_url: str
+    local_mappings: Dict[NormalizedName, MappingEntry]
+
+    def __init__(self) -> None:
+        self.mapping_url = DEFAULT_MAPPING_URL
+        self.local_mappings = {}
 
     @cached_property
-    def pypi_lookup(self) -> Dict[NormalizedName, MappingEntry]:
-        url = self.mapping_url
-        if url.startswith("http://") or url.startswith("https://"):
-            res = requests.get(self._mapping_url)
-            res.raise_for_status()
-            content = res.content
-        else:
-            if url.startswith("file://"):
-                path = url[len("file://") :]
-            else:
-                path = url
-            content = Path(path).read_bytes()
-        lookup = yaml.safe_load(content)
+    def remote_mappings(self) -> Dict[NormalizedName, MappingEntry]:
+        """PyPI to conda name mapping fetched from `mapping_url`"""
+        res = requests.get(self.mapping_url)
+        res.raise_for_status()
+        lookup = yaml.safe_load(res.content)
         # lowercase and kebabcase the pypi names
         assert lookup is not None
         lookup = {canonicalize_name(k): v for k, v in lookup.items()}
         for v in lookup.values():
             v["pypi_name"] = canonicalize_name(v["pypi_name"])
         return lookup
 
+    @property
+    def pypi_lookup(self) -> Dict[NormalizedName, MappingEntry]:
+        """Dict of PyPI to conda name mappings.
+
+        Local mappings take precedence over remote mappings fetched from `mapping_url`.
+        """
+        return {**self.remote_mappings, **self.local_mappings}
+
     @cached_property
     def conda_lookup(self) -> Dict[str, MappingEntry]:
         return {record["conda_name"]: record for record in self.pypi_lookup.values()}
 
 
-LOOKUP_OBJECT = _LookupLoader()
-
+_lookup_loader = _LookupLoader()
 
-def get_forward_lookup() -> Dict[NormalizedName, MappingEntry]:
-    global LOOKUP_OBJECT
-    return LOOKUP_OBJECT.pypi_lookup
 
+def set_lookup_location(lookup_url: str) -> None:
+    """Set the location of the pypi lookup
 
-def get_lookup() -> Dict[str, MappingEntry]:
-    """
-    Reverse grayskull name mapping to map conda names onto PyPI
+    Used by the `lock` cli command to override the DEFAULT_MAPPING_URL for the lookup.
     """
-    global LOOKUP_OBJECT
-    return LOOKUP_OBJECT.conda_lookup
-
-
-def set_lookup_location(lookup_url: str) -> None:
-    global LOOKUP_OBJECT
-    LOOKUP_OBJECT.mapping_url = lookup_url
+    # these will raise AttributeError if they haven't been cached yet.
+    with suppress(AttributeError):
+        del _lookup_loader.remote_mappings
+    with suppress(AttributeError):
+        del _lookup_loader.conda_lookup
+    _lookup_loader.mapping_url = lookup_url
+
+
+def set_pypi_lookup_overrides(mappings: Dict[str, Union[str, MappingEntry]]) -> None:
+    """Set overrides to the pypi lookup"""
+    lookup: Dict[NormalizedName, MappingEntry] = {}
+    # normalize to Dict[NormalizedName, MappingEntry]
+    for k, v in mappings.items():
+        key = canonicalize_name(k)
+        if isinstance(v, dict):
+            if "conda_name" not in v or "pypi_name" not in v:
+                raise ValueError(
+                    "MappingEntries must have both a 'conda_name' and 'pypi_name'"
+                )
+            entry = cast("MappingEntry", dict(v))
+            entry["pypi_name"] = canonicalize_name(str(entry["pypi_name"]))
+        elif isinstance(v, str):
+            entry = {"conda_name": v, "pypi_name": key}
+        else:
+            raise TypeError("Each entry in the mapping must be a string or a dict")
+        lookup[key] = entry
+    _lookup_loader.local_mappings = lookup
 
 
 def conda_name_to_pypi_name(name: str) -> NormalizedName:
     """return the pypi name for a conda package"""
-    lookup = get_lookup()
+    lookup = _lookup_loader.conda_lookup
     cname = canonicalize_name(name)
     return lookup.get(cname, {"pypi_name": cname})["pypi_name"]
 
 
 def pypi_name_to_conda_name(name: str) -> str:
     """return the conda name for a pypi package"""
     cname = canonicalize_name(name)
-    return get_forward_lookup().get(cname, {"conda_name": cname})["conda_name"]
+    forward_lookup = _lookup_loader.pypi_lookup
+    if cname not in forward_lookup:
+        logging.warning(f"Could not find conda name for {cname!r}. Assuming identity.")
+        return cname
+    return forward_lookup[cname]["conda_name"]

conda_lock/src_parser/pyproject_toml.py

@@ -30,7 +30,7 @@
 from typing_extensions import Literal
 
 from conda_lock.common import get_in
-from conda_lock.lookup import get_forward_lookup as get_lookup
+from conda_lock.lookup import pypi_name_to_conda_name, set_pypi_lookup_overrides
 from conda_lock.models.lock_spec import (
     Dependency,
     LockSpecification,
@@ -73,22 +73,6 @@ def join_version_components(pieces: Sequence[Union[str, int]]) -> str:
     return ".".join(str(p) for p in pieces)
 
 
-def normalize_pypi_name(name: str) -> str:
-    cname = canonicalize_pypi_name(name)
-    if cname in get_lookup():
-        lookup = get_lookup()[cname]
-        res = lookup.get("conda_name") or lookup.get("conda_forge")
-        if res is not None:
-            return res
-        else:
-            logging.warning(
-                f"Could not find conda name for {cname}. Assuming identity."
-            )
-            return cname
-    else:
-        return cname
-
-
 def poetry_version_to_conda_version(version_string: Optional[str]) -> Optional[str]:
     if version_string is None:
         return None
@@ -275,7 +259,7 @@ def parse_poetry_pyproject_toml(
                 )
 
             if manager == "conda":
-                name = normalize_pypi_name(depname)
+                name = pypi_name_to_conda_name(depname)
                 version = poetry_version_to_conda_version(poetry_version_spec)
             else:
                 name = depname
@@ -421,16 +405,15 @@ def parse_python_requirement(
 ) -> Dependency:
     """Parse a requirements.txt like requirement to a conda spec"""
     parsed_req = parse_requirement_specifier(requirement)
-    name = canonicalize_pypi_name(parsed_req.name)
     collapsed_version = str(parsed_req.specifier)
     conda_version = poetry_version_to_conda_version(collapsed_version)
     if conda_version:
         conda_version = ",".join(sorted(conda_version.split(",")))
 
     if normalize_name:
-        conda_dep_name = normalize_pypi_name(name)
+        conda_dep_name = pypi_name_to_conda_name(parsed_req.name)
     else:
-        conda_dep_name = name
+        conda_dep_name = canonicalize_pypi_name(parsed_req.name)
     extras = list(parsed_req.extras)
 
     if parsed_req.url and parsed_req.url.startswith("git+"):
@@ -559,6 +542,10 @@ def parse_pyproject_toml(
         contents = toml_load(fp)
     build_system = get_in(["build-system", "build-backend"], contents)
 
+    pypi_map = get_in(["tool", "conda-lock", "pypi-to-conda-name"], contents, False)
+    if pypi_map:
+        set_pypi_lookup_overrides(pypi_map)
+
     if get_in(
         ["tool", "conda-lock", "skip-non-conda-lock"],
         contents,

tests/test-pep621-pypi-override/pyproject.toml

@@ -0,0 +1,10 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "conda-lock-test-pypi-naming"
+dependencies = ["some-name-i-want-to-override"]
+
+[tool.conda-lock.pypi-to-conda-name]
+some-name-i-want-to-override = "resolved-name"

tests/test_conda_lock.py

@@ -347,6 +347,13 @@ def include_dev_dependencies(request: Any) -> bool:
     return request.param
 
 
+@pytest.fixture
+def pep621_pyproject_toml_pypi_override(tmp_path: Path):
+    return clone_test_dir("test-pep621-pypi-override", tmp_path).joinpath(
+        "pyproject.toml"
+    )
+
+
 JSON_FIELDS: Dict[str, str] = {"json_unique_field": "test1", "common_field": "test2"}
 
 YAML_FIELDS: Dict[str, str] = {"yaml_unique_field": "test3", "common_field": "test4"}
@@ -1009,6 +1016,16 @@ def test_parse_poetry_invalid_optionals(pyproject_optional_toml: Path):
     )
 
 
+def test_parse_pyproject_pypi_overrides(pep621_pyproject_toml_pypi_override: Path):
+    res = parse_pyproject_toml(pep621_pyproject_toml_pypi_override, ["linux-64"])
+
+    specs = {dep.name for dep in res.dependencies["linux-64"]}
+
+    # in the pyproject.toml, the package "resolved-name' is provided as an
+    # override for the package "some-name-i-want-to-override".
+    assert "resolved-name" in specs
+
+
 def test_explicit_toposorted() -> None:
     """Verify that explicit lockfiles are topologically sorted.