Support Multiple Categories for Sub-Dependencies in Lockfile (Rebase #390)

.gitignore

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

conda_lock/conda_lock.py

@@ -602,7 +602,7 @@ def render_lockfile_for_platform(  # noqa: C901
         f"# input_hash: {lockfile.metadata.content_hash.get(platform)}\n",
     ]
 
-    categories = {
+    categories_to_install: Set[str] = {
         "main",
         *(extras or []),
         *(["dev"] if include_dev_dependencies else []),
@@ -620,7 +620,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,9 @@ 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] = {
+                category for p in content.package for category in p.categories
+            }
 
             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,12 +1,13 @@
 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,
     DependencySource,
     GitMeta,
     HashModel,
     InputMeta,
+    LockKey,
     LockMeta,
     MetadataOption,
     TimeMeta,
@@ -17,20 +18,32 @@
 
 
 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]:
+        """Convert a v2 dependency into a list of v1 dependencies.
+
+        In case a v2 dependency might contain multiple categories, but a v1 dependency
+        can only contain a single category, we represent multiple categories as a list
+        of v1 dependencies that are identical except for the `category` field. The
+        `category` field runs over all categories."""
+        package_entries_per_category = [
+            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)
+        ]
+        return package_entries_per_category
 
 
 class Lockfile(StrictModel):
@@ -127,35 +140,71 @@ def _toposort(package: List[LockedDependency]) -> List[LockedDependency]:
         return final_package
 
     def to_v1(self) -> LockfileV1:
+        # Each v2 package gives a list of v1 packages.
+        # Flatten these into a single list of v1 packages.
+        v1_packages = [
+            package_entry_per_category
+            for p in self.package
+            for package_entry_per_category in p.to_v1()
+        ]
         return LockfileV1(
-            package=[p.to_v1() for p in self.package],
+            package=v1_packages,
             metadata=self.metadata,
         )
 
 
-def _locked_dependency_v1_to_v2(dep: LockedDependencyV1) -> LockedDependency:
+def _locked_dependency_v1_to_v2(
+    package_entries_per_category: List[LockedDependencyV1],
+) -> LockedDependency:
     """Convert a LockedDependency from v1 to v2.
 
-    * Remove the optional field (it is always equal to category != "main")
+    This is an inverse to `LockedDependency.to_v1()`.
     """
+    # Dependencies are parsed from a v1 lockfile, so there will always be
+    # at least one entry corresponding to what was parsed.
+    assert len(package_entries_per_category) > 0
+    # All the package entries should share the same key.
+    assert all(
+        d.key() == package_entries_per_category[0].key()
+        for d in package_entries_per_category
+    )
+
+    categories = {d.category for d in package_entries_per_category}
+
+    # Each entry should correspond to a distinct category
+    assert len(categories) == len(package_entries_per_category)
+
     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=package_entries_per_category[0].name,
+        version=package_entries_per_category[0].version,
+        manager=package_entries_per_category[0].manager,
+        platform=package_entries_per_category[0].platform,
+        dependencies=package_entries_per_category[0].dependencies,
+        url=package_entries_per_category[0].url,
+        hash=package_entries_per_category[0].hash,
+        categories=categories,
+        source=package_entries_per_category[0].source,
+        build=package_entries_per_category[0].build,
     )
 
 
 def lockfile_v1_to_v2(lockfile_v1: LockfileV1) -> Lockfile:
-    """Convert a Lockfile from v1 to v2."""
+    """Convert a Lockfile from v1 to v2.
+
+    Entries may share the same key if they represent a dependency
+    belonging to multiple categories. They must be collected here.
+    """
+    dependencies_for_key: Dict[LockKey, List[LockedDependencyV1]] = defaultdict(list)
+    for dep in lockfile_v1.package:
+        dependencies_for_key[dep.key()].append(dep)
+
+    v2_packages = [
+        _locked_dependency_v1_to_v2(package_entries_per_category)
+        for package_entries_per_category in dependencies_for_key.values()
+    ]
+
     return Lockfile(
-        package=[_locked_dependency_v1_to_v2(p) for p in lockfile_v1.package],
+        package=v2_packages,
         metadata=lockfile_v1.metadata,
     )