experiments with categories

.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: Set[str] = {
         "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

@@ -131,7 +131,7 @@ def dep_name(manager: str, dep: str) -> str:
         if not isinstance(targets, list):
             targets = [targets]
         for target in targets:
-            target.category = source.category
+            target.categories = {source.category}
 
 
 def parse_conda_lock_file(path: pathlib.Path) -> Lockfile:
@@ -163,7 +163,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,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] = {"main"}
+
+    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):
@@ -120,35 +133,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,
     )
 

tests/test_conda_lock.py

@@ -12,7 +12,6 @@
 import tempfile
 import typing
 import uuid
-import warnings
 
 from glob import glob
 from pathlib import Path
@@ -385,10 +384,10 @@ def test_lock_poetry_ibis(
     )
     lockfile = parse_conda_lock_file(pyproject.parent / DEFAULT_LOCKFILE_NAME)
 
-    all_categories = set()
+    all_categories: Set[str] = set()
 
     for pkg in lockfile.package:
-        all_categories.add(pkg.category)
+        all_categories.update(pkg.categories)
 
     for desired_category in extra_categories:
         assert (