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

README.md

@@ -331,6 +331,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`:
+
+```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,5 +1,8 @@
+import logging
+
+from contextlib import suppress
 from functools import cached_property
-from typing import Dict
+from typing import ClassVar, Dict, Optional, Union, cast
 
 import requests
 import yaml
@@ -8,28 +11,51 @@
 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."""
+
+    _SINGLETON: ClassVar[Optional["_LookupLoader"]] = None
+
+    @classmethod
+    def instance(cls) -> "_LookupLoader":
+        if cls._SINGLETON is None:
+            cls._SINGLETON = cls()
+        return cls._SINGLETON
+
+    def __init__(
+        self,
+        pypi_lookup_overrides: Optional[Dict[NormalizedName, MappingEntry]] = None,
+        mapping_url: str = DEFAULT_MAPPING_URL,
+    ) -> None:
+        self._mapping_url = mapping_url
+        self._local_mappings = pypi_lookup_overrides
 
     @property
     def mapping_url(self) -> str:
         return self._mapping_url
 
     @mapping_url.setter
     def mapping_url(self, value: str) -> None:
-        del self.pypi_lookup
-        del self.conda_lookup
+        # these will raise AttributeError if they haven't been cached yet.
+        with suppress(AttributeError):
+            del self.remote_mappings
+        with suppress(AttributeError):
+            del self.pypi_lookup
+        with suppress(AttributeError):
+            del self.conda_lookup
         self._mapping_url = value
 
     @cached_property
-    def pypi_lookup(self) -> Dict[NormalizedName, MappingEntry]:
+    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)
@@ -40,40 +66,71 @@ def pypi_lookup(self) -> Dict[NormalizedName, MappingEntry]:
             v["pypi_name"] = canonicalize_name(v["pypi_name"])
         return lookup
 
-    @cached_property
-    def conda_lookup(self) -> Dict[str, MappingEntry]:
-        return {record["conda_name"]: record for record in self.pypi_lookup.values()}
+    @property
+    def local_mappings(self) -> Dict[NormalizedName, MappingEntry]:
+        """PyPI to conda name mappings set by the user."""
+        return self._local_mappings or {}
+
+    @local_mappings.setter
+    def local_mappings(self, mappings: Dict[str, Union[str, MappingEntry]]) -> None:
+        """Value should be a mapping from pypi name to conda name or a mapping entry."""
+        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
+        self._local_mappings = lookup
 
+    @cached_property
+    def pypi_lookup(self) -> Dict[NormalizedName, MappingEntry]:
+        """Dict of PyPI to conda name mappings.
 
-LOOKUP_OBJECT = _LookupLoader()
+        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()}
 
-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]:
+    Used by the `lock` cli command to override the DEFAULT_MAPPING_URL for the lookup.
     """
-    Reverse grayskull name mapping to map conda names onto PyPI
-    """
-    global LOOKUP_OBJECT
-    return LOOKUP_OBJECT.conda_lookup
+    _LookupLoader.instance().mapping_url = lookup_url
 
 
-def set_lookup_location(lookup_url: str) -> None:
-    global LOOKUP_OBJECT
-    LOOKUP_OBJECT.mapping_url = lookup_url
+def set_pypi_lookup_overrides(mappings: Dict[str, Union[str, MappingEntry]]) -> None:
+    """Set overrides to the pypi lookup"""
+    # type ignore because the setter will normalize the types
+    _LookupLoader.instance().local_mappings = mappings  # type: ignore [assignment]
 
 
 def conda_name_to_pypi_name(name: str) -> NormalizedName:
     """return the pypi name for a conda package"""
-    lookup = get_lookup()
+    lookup = _LookupLoader.instance().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 = _LookupLoader.instance().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,