Attempted rebase of #390

conda_lock/conda_lock.py

@@ -542,17 +542,15 @@ def render_lockfile_for_platform(  # noqa: C901
     lockfile.toposort_inplace()
 
     for p in lockfile.package:
-        if p.platform == platform and p.category in categories:
+        if p.platform == platform and not p.categories.isdisjoint(categories):
             if p.manager == "pip":
                 pip_deps.append(p)
             elif p.manager == "conda":
                 # exclude virtual packages
                 if not p.name.startswith("__"):
                     conda_deps.append(p)
 
-    def format_pip_requirement(
-        spec: LockedDependency, platform: str, direct: bool = False
-    ) -> str:
+    def format_pip_requirement(spec: LockedDependency, direct: bool = False) -> str:
         if spec.source and spec.source.type == "url":
             return f"{spec.name} @ {spec.source.url}"
         elif direct:
@@ -566,9 +564,7 @@ def format_pip_requirement(
                 s += f" --hash=sha256:{spec.hash.sha256}"
             return s
 
-    def format_conda_requirement(
-        spec: LockedDependency, platform: str, direct: bool = False
-    ) -> str:
+    def format_conda_requirement(spec: LockedDependency, direct: bool = False) -> str:
         if direct:
             # inject the environment variables in here
             return posixpath.expandvars(f"{spec.url}#{spec.hash.md5}")
@@ -589,7 +585,7 @@ def format_conda_requirement(
                 ),
                 "dependencies:",
                 *(
-                    f"  - {format_conda_requirement(dep, platform, direct=False)}"
+                    f"  - {format_conda_requirement(dep, direct=False)}"
                     for dep in conda_deps
                 ),
             ]
@@ -598,7 +594,7 @@ def format_conda_requirement(
             [
                 "  - pip:",
                 *(
-                    f"    - {format_pip_requirement(dep, platform, direct=False)}"
+                    f"    - {format_pip_requirement(dep, direct=False)}"
                     for dep in pip_deps
                 ),
             ]
@@ -609,7 +605,7 @@ def format_conda_requirement(
         lockfile_contents.append("@EXPLICIT\n")
 
         lockfile_contents.extend(
-            [format_conda_requirement(dep, platform, direct=True) for dep in conda_deps]
+            [format_conda_requirement(dep, direct=True) for dep in conda_deps]
         )
 
         def sanitize_lockfile_line(line: str) -> str:
@@ -623,10 +619,7 @@ def sanitize_lockfile_line(line: str) -> str:
 
         # emit an explicit requirements.txt, prefixed with '# pip '
         lockfile_contents.extend(
-            [
-                f"# pip {format_pip_requirement(dep, platform, direct=True)}"
-                for dep in pip_deps
-            ]
+            [f"# pip {format_pip_requirement(dep, direct=True)}" for dep in pip_deps]
         )
 
         if len(pip_deps) > 0:

conda_lock/lockfile/__init__.py

@@ -1,9 +1,18 @@
-import json
 import pathlib
 
 from collections import defaultdict
 from textwrap import dedent
-from typing import Any, Collection, Dict, List, Mapping, Optional, Sequence, Set, Union
+from typing import (
+    Any,
+    Collection,
+    DefaultDict,
+    Dict,
+    List,
+    Mapping,
+    Optional,
+    Set,
+    Union,
+)
 
 import yaml
 
@@ -38,7 +47,6 @@ def _seperator_munge_get(
 def apply_categories(
     requested: Dict[str, Dependency],
     planned: Mapping[str, Union[List[LockedDependency], LockedDependency]],
-    categories: Sequence[str] = ("main", "dev"),
     convert_to_pip_names: bool = False,
 ) -> None:
     """map each package onto the root request the with the highest-priority category"""
@@ -57,7 +65,6 @@ def apply_categories(
 
     # walk dependency tree to assemble all transitive dependencies by request
     dependents: Dict[str, Set[str]] = {}
-    by_category = defaultdict(list)
 
     def extract_planned_items(
         planned_items: Union[List[LockedDependency], LockedDependency]
@@ -78,8 +85,8 @@ def dep_name(manager: str, dep: str) -> str:
             return conda_name_to_pypi_name(dep).lower()
         return dep
 
-    for name, request in requested.items():
-        todo: List[str] = list()
+    for name in requested:
+        todo: List[str] = []
         deps: Set[str] = set()
         item = name
 
@@ -106,29 +113,22 @@ def dep_name(manager: str, dep: str) -> str:
 
         dependents[name] = deps
 
-        by_category[request.category].append(request.name)
-
-    # now, map each package to its root request preferring the ones earlier in the
-    # list
-    categories = [*categories, *(k for k in by_category if k not in categories)]
-    root_requests = {}
-    for category in categories:
-        for root in by_category.get(category, []):
-            for transitive_dep in dependents[root]:
-                if transitive_dep not in root_requests:
-                    root_requests[transitive_dep] = root
+    # now, map each package to its root requests / dependencies
+    root_requests: DefaultDict[str, Set[str]] = defaultdict(set)
+    for root, transitive_deps in dependents.items():
+        for transitive_dep in transitive_deps:
+            root_requests[transitive_dep].add(root)
+
     # include root requests themselves
     for name in requested:
-        root_requests[name] = name
+        root_requests[name].add(name)
 
-    for dep, root in root_requests.items():
-        source = requested[root]
-        # try a conda target first
-        targets = _seperator_munge_get(planned, dep)
-        if not isinstance(targets, list):
-            targets = [targets]
-        for target in targets:
-            target.category = source.category
+    for dep, roots in root_requests.items():
+        target = _seperator_munge_get(planned, dep)
+        for root in roots:
+            source = requested[root]
+            assert isinstance(target, LockedDependency)  # TODO: why?
+            target.categories.add(source.category)
 
 
 def parse_conda_lock_file(path: pathlib.Path) -> Lockfile:
@@ -141,10 +141,13 @@ def parse_conda_lock_file(path: pathlib.Path) -> Lockfile:
     if not (isinstance(version, int) and version <= Lockfile.version):
         raise ValueError(f"{path} has unknown version {version}")
 
+    packages = {}
     for p in content["package"]:
+        del p["category"]
         del p["optional"]
+        packages[(p["name"], p["version"], p["platform"])] = p
 
-    return Lockfile.parse_obj(content)
+    return Lockfile.parse_obj({**content, "package": list(packages.values())})
 
 
 def write_conda_lock_file(
@@ -156,7 +159,7 @@ def write_conda_lock_file(
     content.toposort_inplace()
     with path.open("w") as f:
         if include_help_text:
-            categories = set(p.category for p in content.package)
+            categories = {cat for p in content.package for cat in p.categories}
 
             def write_section(text: str) -> None:
                 lines = dedent(text).split("\n")
@@ -209,20 +212,28 @@ def write_section(text: str) -> None:
 
         output: Dict[str, Any] = {
             "version": Lockfile.version,
-            "metadata": json.loads(
-                content.metadata.json(
-                    by_alias=True, exclude_unset=True, exclude_none=True
-                )
+            "metadata": content.metadata.dict(
+                by_alias=True, exclude_unset=True, exclude_none=True
             ),
-            "package": [
-                {
-                    **package.dict(
-                        by_alias=True, exclude_unset=True, exclude_none=True
-                    ),
-                    "optional": (package.category != "main"),
-                }
-                for package in content.package
-            ],
+            "package": [],
         }
 
+        for package in content.package:
+            sorted_cats = sorted(package.categories)
+            for category in sorted_cats:
+                output["package"].append(
+                    dict(
+                        sorted(
+                            {
+                                **package.dict(
+                                    by_alias=True, exclude_unset=True, exclude_none=True
+                                ),
+                                "categories": sorted_cats,
+                                "category": category,
+                                "optional": (category != "main"),
+                            }.items()
+                        )
+                    )
+                )
+
         yaml.dump(output, stream=f, sort_keys=False)

conda_lock/lockfile/models.py

@@ -6,7 +6,16 @@
 import typing
 
 from collections import defaultdict, namedtuple
-from typing import TYPE_CHECKING, AbstractSet, ClassVar, Dict, List, Optional, Union
+from typing import (
+    TYPE_CHECKING,
+    AbstractSet,
+    ClassVar,
+    Dict,
+    List,
+    Optional,
+    Set,
+    Union,
+)
 
 
 if TYPE_CHECKING:
@@ -44,7 +53,7 @@ class LockedDependency(StrictModel):
     dependencies: Dict[str, str] = {}
     url: str
     hash: HashModel
-    category: str = "main"
+    categories: Set[str] = set()
     source: Optional[DependencySource] = None
     build: Optional[str] = None
 

tests/gdal/environment.yml

@@ -6,4 +6,4 @@ dependencies:
   - python >= 3.7, < 3.8
   - gdal
   - pip:
-    - toolz
+    - toolz

tests/test_conda_lock.py

@@ -33,6 +33,7 @@
     _add_auth_to_line,
     _add_auth_to_lockfile,
     _extract_domain,
+    _solve_for_arch,
     _strip_auth_from_line,
     _strip_auth_from_lockfile,
     create_lockfile_from_spec,
@@ -346,7 +347,7 @@ def test_lock_poetry_ibis(
     all_categories = set()
 
     for pkg in lockfile.package:
-        all_categories.add(pkg.category)
+        all_categories.update(pkg.categories)
 
     for desired_category in extra_categories:
         assert (
@@ -520,6 +521,7 @@ def test_choose_wheel() -> None:
         platform="linux-64",
     )
     assert len(solution) == 1
+    assert solution["fastavro"].categories == {"main"}
     assert solution["fastavro"].hash == HashModel(
         sha256="a111a384a786b7f1fd6a8a8307da07ccf4d4c425084e2d61bae33ecfb60de405"
     )
@@ -1183,6 +1185,54 @@ def test_run_lock_with_input_hash_check(
     assert "Spec hash already locked for" in output.err
 
 
+def test_solve_arch_multiple_categories():
+    _conda_exe = determine_conda_executable(None, mamba=False, micromamba=False)
+    vpr = default_virtual_package_repodata()
+    channels = [Channel.from_string("conda-forge")]
+
+    with vpr, tempfile.NamedTemporaryFile(dir=".") as tf:
+        spec = LockSpecification(
+            dependencies={
+                "linux-64": [
+                    VersionedDependency(
+                        name="python",
+                        version="=3.10.9",
+                        manager="conda",
+                        category="main",
+                        extras=[],
+                    ),
+                    VersionedDependency(
+                        name="pandas",
+                        version="=1.5.3",
+                        manager="conda",
+                        category="test",
+                        extras=[],
+                    ),
+                    VersionedDependency(
+                        name="pyarrow",
+                        version="=9.0.0",
+                        manager="conda",
+                        category="dev",
+                        extras=[],
+                    ),
+                ],
+            },
+            channels=channels,
+            # NB: this file must exist for relative path resolution to work
+            # in create_lockfile_from_spec
+            sources=[Path(tf.name)],
+            virtual_package_repo=vpr,
+        )
+
+        locked_deps = _solve_for_arch(_conda_exe, spec, "linux-64", channels)
+        python_deps = [dep for dep in locked_deps if dep.name == "python"]
+        assert len(python_deps) == 1
+        assert python_deps[0].categories == {"main", "test", "dev"}
+        numpy_deps = [dep for dep in locked_deps if dep.name == "numpy"]
+        assert len(numpy_deps) == 1
+        assert numpy_deps[0].categories == {"test", "dev"}
+
+
 @pytest.mark.parametrize(
     "package,version,url_pattern",
     [