Lockfile-Related Code Refactored Before Sourcefile PR

conda_lock/conda_lock.py

@@ -62,22 +62,22 @@
     PIP_SUPPORT = True
 except ImportError:
     PIP_SUPPORT = False
-from conda_lock.lookup import set_lookup_location
-from conda_lock.src_parser import (
+from conda_lock.lockfile import (
     Dependency,
     GitMeta,
     InputMeta,
     LockedDependency,
     Lockfile,
     LockMeta,
-    LockSpecification,
     MetadataOption,
     TimeMeta,
     UpdateSpecification,
-    aggregate_lock_specs,
+    parse_conda_lock_file,
+    write_conda_lock_file,
 )
+from conda_lock.lookup import set_lookup_location
+from conda_lock.src_parser import LockSpecification, aggregate_lock_specs
 from conda_lock.src_parser.environment_yaml import parse_environment_file
-from conda_lock.src_parser.lockfile import parse_conda_lock_file, write_conda_lock_file
 from conda_lock.src_parser.meta_yaml import parse_meta_yaml_file
 from conda_lock.src_parser.pyproject_toml import parse_pyproject_toml
 from conda_lock.virtual_package import (

conda_lock/conda_solver.py

@@ -31,14 +31,9 @@
     conda_pkgs_dir,
     is_micromamba,
 )
+from conda_lock.lockfile import HashModel, LockedDependency, _apply_categories
 from conda_lock.models.channel import Channel
-from conda_lock.src_parser import (
-    Dependency,
-    HashModel,
-    LockedDependency,
-    VersionedDependency,
-    _apply_categories,
-)
+from conda_lock.src_parser import Dependency, VersionedDependency
 
 
 logger = logging.getLogger(__name__)

conda_lock/src_parser/lockfile.py → conda_lock/lockfile/__init__.py

@@ -1,14 +1,87 @@
 import json
 import pathlib
 
+from collections import defaultdict
 from textwrap import dedent
-from typing import Collection, Optional
+from typing import Collection, Dict, List, Optional, Sequence, Set
 
 import yaml
 
-from conda_lock.src_parser import MetadataOption
+from conda_lock.src_parser import Dependency
 
-from . import Lockfile
+from .models import DependencySource as DependencySource
+from .models import GitMeta as GitMeta
+from .models import HashModel as HashModel
+from .models import InputMeta as InputMeta
+from .models import LockedDependency, Lockfile
+from .models import LockKey as LockKey
+from .models import LockMeta as LockMeta
+from .models import MetadataOption
+from .models import TimeMeta as TimeMeta
+from .models import UpdateSpecification as UpdateSpecification
+
+
+def _apply_categories(
+    requested: Dict[str, Dependency],
+    planned: Dict[str, LockedDependency],
+    categories: Sequence[str] = ("main", "dev"),
+) -> None:
+    """map each package onto the root request the with the highest-priority category"""
+    # walk dependency tree to assemble all transitive dependencies by request
+    dependents: Dict[str, Set[str]] = {}
+    by_category = defaultdict(list)
+
+    def seperator_munge_get(
+        d: Dict[str, LockedDependency], key: str
+    ) -> LockedDependency:
+        # since separators are not consistent across managers (or even within) we need to do some double attempts here
+        try:
+            return d[key]
+        except KeyError:
+            try:
+                return d[key.replace("-", "_")]
+            except KeyError:
+                return d[key.replace("_", "-")]
+
+    for name, request in requested.items():
+        todo: List[str] = list()
+        deps: Set[str] = set()
+        item = name
+        while True:
+            todo.extend(
+                dep
+                for dep in seperator_munge_get(planned, item).dependencies
+                # exclude virtual packages
+                if not (dep in deps or dep.startswith("__"))
+            )
+            if todo:
+                item = todo.pop(0)
+                deps.add(item)
+            else:
+                break
+
+        dependents[name] = deps
+
+        by_category[request.category].append(request.name)
+
+    # now, map each package to its root request
+    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
+    # include root requests themselves
+    for name in requested:
+        root_requests[name] = name
+
+    for dep, root in root_requests.items():
+        source = requested[root]
+        # try a conda target first
+        target = seperator_munge_get(planned, dep)
+        target.category = source.category
+        target.optional = source.optional
 
 
 def parse_conda_lock_file(