Support Multiple Categories for Sub-Dependencies in Lockfile

.gitignore

@@ -1,3 +1,4 @@
+**.DS_Store
 *.egg-info
 *.eggs
 *.pyc

conda_lock/conda_lock.py

@@ -598,7 +598,7 @@ def render_lockfile_for_platform(  # noqa: C901
         f"# input_hash: {lockfile.metadata.content_hash.get(platform)}\n",
     ]
 
-    categories = {
+    categories_to_install = {
         "main",
         *(extras or []),
         *(["dev"] if include_dev_dependencies else []),
@@ -616,7 +616,7 @@ def render_lockfile_for_platform(  # noqa: C901
     lockfile.filter_virtual_packages_inplace()
 
     for p in lockfile.package:
-        if p.platform == platform and p.category in categories:
+        if p.platform == platform and len(p.categories & categories_to_install) > 0:
             if p.manager == "pip":
                 pip_deps.append(p)
             elif p.manager == "conda":

conda_lock/lockfile/__init__.py

@@ -2,7 +2,17 @@
 
 from collections import defaultdict
 from textwrap import dedent
-from typing import Collection, Dict, List, Mapping, Optional, Sequence, Set, Union
+from typing import (
+    Collection,
+    DefaultDict,
+    Dict,
+    List,
+    Mapping,
+    Optional,
+    Sequence,
+    Set,
+    Union,
+)
 
 import yaml
 
@@ -38,6 +48,23 @@ def _seperator_munge_get(
             return d[key.replace("_", "-")]
 
 
+def _truncate_main_category(
+    planned: Mapping[str, Union[List[LockedDependency], LockedDependency]],
+) -> None:
+    """
+    Given the package dependencies with their respective categories
+    for any package that is in the main category, remove all other associated categories
+    """
+    # Packages in the main category are always installed
+    # so other categories are not necessary
+    for targets in planned.values():
+        if not isinstance(targets, list):
+            targets = [targets]
+        for target in targets:
+            if "main" in target.categories:
+                target.categories = {"main"}
+
+
 def apply_categories(
     requested: Dict[str, Dependency],
     planned: Mapping[str, Union[List[LockedDependency], LockedDependency]],
@@ -111,27 +138,31 @@ def dep_name(manager: str, dep: str) -> str:
 
         by_category[request.category].append(request.name)
 
-    # now, map each package to its root request preferring the ones earlier in the
-    # list
+    # now, map each package to every root request that requires it
     categories = [*categories, *(k for k in by_category if k not in categories)]
-    root_requests = {}
+    root_requests: DefaultDict[str, List[str]] = defaultdict(list)
     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
+                root_requests[transitive_dep].append(root)
     # include root requests themselves
     for name in requested:
-        root_requests[name] = name
+        root_requests[name].append(name)
 
-    for dep, root in root_requests.items():
-        source = requested[root]
+    for dep, roots in root_requests.items():
         # 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 root in roots:
+            source = requested[root]
+            for target in targets:
+                target.categories.add(source.category)
+
+    # For any dep that is part of the 'main' category
+    # we should remove all other categories
+    _truncate_main_category(planned)
 
 
 def parse_conda_lock_file(path: pathlib.Path) -> Lockfile:
@@ -163,7 +194,7 @@ def write_conda_lock_file(
     content.filter_virtual_packages_inplace()
     with path.open("w") as f:
         if include_help_text:
-            categories = set(p.category for p in content.package)
+            categories: Set[str] = set().union(*(p.categories for p in content.package))
 
             def write_section(text: str) -> None:
                 lines = dedent(text).split("\n")

conda_lock/lockfile/v1/models.py

@@ -54,7 +54,6 @@ class BaseLockedDependency(StrictModel):
     dependencies: Dict[str, str] = {}
     url: str
     hash: HashModel
-    category: str = "main"
     source: Optional[DependencySource] = None
     build: Optional[str] = None
 
@@ -69,6 +68,7 @@ def validate_hash(cls, v: HashModel, values: Dict[str, typing.Any]) -> HashModel
 
 
 class LockedDependency(BaseLockedDependency):
+    category: str = "main"
     optional: bool
 
 

conda_lock/lockfile/v2prelim/models.py

@@ -1,5 +1,5 @@
 from collections import defaultdict
-from typing import ClassVar, Dict, List, Optional
+from typing import ClassVar, Dict, List, Optional, Set
 
 from conda_lock.lockfile.v1.models import (
     BaseLockedDependency,
@@ -17,20 +17,25 @@
 
 
 class LockedDependency(BaseLockedDependency):
-    def to_v1(self) -> LockedDependencyV1:
-        return LockedDependencyV1(
-            name=self.name,
-            version=self.version,
-            manager=self.manager,
-            platform=self.platform,
-            dependencies=self.dependencies,
-            url=self.url,
-            hash=self.hash,
-            category=self.category,
-            source=self.source,
-            build=self.build,
-            optional=self.category != "main",
-        )
+    categories: Set[str] = set()
+
+    def to_v1(self) -> List[LockedDependencyV1]:
+        return [
+            LockedDependencyV1(
+                name=self.name,
+                version=self.version,
+                manager=self.manager,
+                platform=self.platform,
+                dependencies=self.dependencies,
+                url=self.url,
+                hash=self.hash,
+                category=category,
+                source=self.source,
+                build=self.build,
+                optional=category != "main",
+            )
+            for category in sorted(self.categories)
+        ]
 
 
 class Lockfile(StrictModel):
@@ -121,34 +126,45 @@ def _toposort(package: List[LockedDependency]) -> List[LockedDependency]:
 
     def to_v1(self) -> LockfileV1:
         return LockfileV1(
-            package=[p.to_v1() for p in self.package],
+            package=[out for p in self.package for out in p.to_v1()],
             metadata=self.metadata,
         )
 
 
-def _locked_dependency_v1_to_v2(dep: LockedDependencyV1) -> LockedDependency:
+def _locked_dependency_v1_to_v2(dep: List[LockedDependencyV1]) -> LockedDependency:
     """Convert a LockedDependency from v1 to v2.
 
     * Remove the optional field (it is always equal to category != "main")
     """
+    assert len(dep) > 0
+    assert all(d.key() == dep[0].key() for d in dep)
+    assert len(set(d.category for d in dep)) == len(dep)
+
     return LockedDependency(
-        name=dep.name,
-        version=dep.version,
-        manager=dep.manager,
-        platform=dep.platform,
-        dependencies=dep.dependencies,
-        url=dep.url,
-        hash=dep.hash,
-        category=dep.category,
-        source=dep.source,
-        build=dep.build,
+        name=dep[0].name,
+        version=dep[0].version,
+        manager=dep[0].manager,
+        platform=dep[0].platform,
+        dependencies=dep[0].dependencies,
+        url=dep[0].url,
+        hash=dep[0].hash,
+        categories={d.category for d in dep},
+        source=dep[0].source,
+        build=dep[0].build,
     )
 
 
 def lockfile_v1_to_v2(lockfile_v1: LockfileV1) -> Lockfile:
     """Convert a Lockfile from v1 to v2."""
+    final_dependencies = defaultdict(list)
+    for dep in lockfile_v1.package:
+        final_dependencies[dep.key()].append(dep)
+
     return Lockfile(
-        package=[_locked_dependency_v1_to_v2(p) for p in lockfile_v1.package],
+        package=[
+            _locked_dependency_v1_to_v2(v1_pkgs)
+            for v1_pkgs in final_dependencies.values()
+        ],
         metadata=lockfile_v1.metadata,
     )