Refactor LockSpecification as a Dictionary from Platforms to List of Deps

conda_lock/conda_lock.py

@@ -656,12 +656,8 @@ def _solve_for_arch(
     """
     if update_spec is None:
         update_spec = UpdateSpecification()
-    # filter requested and locked dependencies to the current platform
-    dependencies = [
-        dep
-        for dep in spec.dependencies
-        if (not dep.selectors.platform) or platform in dep.selectors.platform
-    ]
+
+    dependencies = spec.dependencies[platform]
     locked = [dep for dep in update_spec.locked if dep.platform == platform]
     requested_deps_by_name = {
         manager: {dep.name: dep for dep in dependencies if dep.manager == manager}

conda_lock/conda_solver.py

@@ -9,17 +9,7 @@
 import time
 
 from contextlib import contextmanager
-from typing import (
-    Any,
-    Dict,
-    Iterable,
-    Iterator,
-    List,
-    MutableSequence,
-    Optional,
-    Sequence,
-    cast,
-)
+from typing import Dict, Iterable, Iterator, List, MutableSequence, Optional, Sequence
 from urllib.parse import urlsplit, urlunsplit
 
 import yaml

conda_lock/models/lock_spec.py

@@ -13,29 +13,12 @@
 from conda_lock.virtual_package import FakeRepoData
 
 
-class Selectors(StrictModel):
-    platform: Optional[List[str]] = None
-
-    def __ior__(self, other: "Selectors") -> "Selectors":
-        if not isinstance(other, Selectors):
-            raise TypeError
-        if other.platform and self.platform:
-            for p in other.platform:
-                if p not in self.platform:
-                    self.platform.append(p)
-        return self
-
-    def for_platform(self, platform: str) -> bool:
-        return self.platform is None or platform in self.platform
-
-
 class _BaseDependency(StrictModel):
     name: str
     manager: Literal["conda", "pip"] = "conda"
     optional: bool = False
     category: str = "main"
     extras: List[str] = []
-    selectors: Selectors = Selectors()
 
 
 class VersionedDependency(_BaseDependency):
@@ -58,14 +41,17 @@ class Package(StrictModel):
 
 
 class LockSpecification(BaseModel):
-    dependencies: List[Dependency]
+    dependencies: Dict[str, List[Dependency]]
     # TODO: Should we store the auth info in here?
     channels: List[Channel]
-    platforms: List[str]
     sources: List[pathlib.Path]
     virtual_package_repo: Optional[FakeRepoData] = None
     allow_pypi_requests: bool = True
 
+    @property
+    def platforms(self) -> List[str]:
+        return list(self.dependencies.keys())
+
     def content_hash(self) -> Dict[str, str]:
         return {
             platform: self.content_hash_for_platform(platform)
@@ -77,8 +63,9 @@ def content_hash_for_platform(self, platform: str) -> str:
             "channels": [c.json() for c in self.channels],
             "specs": [
                 p.dict()
-                for p in sorted(self.dependencies, key=lambda p: (p.manager, p.name))
-                if p.selectors.for_platform(platform)
+                for p in sorted(
+                    self.dependencies[platform], key=lambda p: (p.manager, p.name)
+                )
             ],
         }
         if self.virtual_package_repo is not None:

conda_lock/src_parser/__init__.py

@@ -12,7 +12,10 @@
     parse_platforms_from_env_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.src_parser.pyproject_toml import (
+    parse_platforms_from_pyproject_toml,
+    parse_pyproject_toml,
+)
 from conda_lock.virtual_package import FakeRepoData
 
 
@@ -36,7 +39,7 @@ def _parse_platforms_from_srcs(src_files: List[pathlib.Path]) -> List[str]:
         if src_file.name == "meta.yaml":
             continue
         elif src_file.name == "pyproject.toml":
-            all_file_platforms.append(parse_pyproject_toml(src_file).platforms)
+            all_file_platforms.append(parse_platforms_from_pyproject_toml(src_file))
         else:
             all_file_platforms.append(parse_platforms_from_env_file(src_file))
 
@@ -62,7 +65,7 @@ def _parse_source_files(
         if src_file.name == "meta.yaml":
             desired_envs.append(parse_meta_yaml_file(src_file, platforms))
         elif src_file.name == "pyproject.toml":
-            desired_envs.append(parse_pyproject_toml(src_file))
+            desired_envs.append(parse_pyproject_toml(src_file, platforms))
         else:
             desired_envs.append(parse_environment_file(src_file, platforms))
     return desired_envs
@@ -87,7 +90,6 @@ def make_lock_spec(
 
     lock_spec = aggregate_lock_specs(lock_specs)
     lock_spec.virtual_package_repo = virtual_package_repo
-    lock_spec.platforms = platforms
     lock_spec.channels = (
         [Channel.from_string(co) for co in channel_overrides]
         if channel_overrides
@@ -99,10 +101,13 @@ def make_lock_spec(
         def dep_has_category(d: Dependency, categories: AbstractSet[str]) -> bool:
             return d.category in categories
 
-        lock_spec.dependencies = [
-            d
-            for d in lock_spec.dependencies
-            if dep_has_category(d, categories=required_categories)
-        ]
+        lock_spec.dependencies = {
+            platform: [
+                d
+                for d in dependencies
+                if dep_has_category(d, categories=required_categories)
+            ]
+            for platform, dependencies in lock_spec.dependencies.items()
+        }
 
     return lock_spec

conda_lock/src_parser/aggregation.py

@@ -14,20 +14,21 @@
 def aggregate_lock_specs(
     lock_specs: List[LockSpecification],
 ) -> LockSpecification:
-    # unique dependencies
-    unique_deps: Dict[Tuple[str, str], Dependency] = {}
-    for dep in chain.from_iterable(
-        [lock_spec.dependencies for lock_spec in lock_specs]
-    ):
-        key = (dep.manager, dep.name)
-        if key in unique_deps:
-            # Override existing, but merge selectors
-            previous_selectors = unique_deps[key].selectors
-            previous_selectors |= dep.selectors
-            dep.selectors = previous_selectors
-        unique_deps[key] = dep
-
-    dependencies = list(unique_deps.values())
+    # Preserve input order of platforms
+    platforms = ordered_union(lock_spec.platforms or [] for lock_spec in lock_specs)
+
+    dependencies: Dict[str, List[Dependency]] = {}
+    for platform in platforms:
+        # unique dependencies
+        unique_deps: Dict[Tuple[str, str], Dependency] = {}
+        for dep in chain.from_iterable(
+            lock_spec.dependencies.get(platform, []) for lock_spec in lock_specs
+        ):
+            key = (dep.manager, dep.name)
+            unique_deps[key] = dep
+
+        dependencies[platform] = list(unique_deps.values())
+
     try:
         channels = suffix_union(lock_spec.channels or [] for lock_spec in lock_specs)
     except ValueError as e:
@@ -38,7 +39,6 @@ def aggregate_lock_specs(
         # Ensure channel are correctly ordered
         channels=channels,
         # uniquify metadata, preserving order
-        platforms=ordered_union(lock_spec.platforms or [] for lock_spec in lock_specs),
         sources=ordered_union(lock_spec.sources or [] for lock_spec in lock_specs),
         allow_pypi_requests=all(
             lock_spec.allow_pypi_requests for lock_spec in lock_specs

conda_lock/src_parser/environment_yaml.py

@@ -7,7 +7,6 @@
 import yaml
 
 from conda_lock.models.lock_spec import Dependency, LockSpecification
-from conda_lock.src_parser.aggregation import aggregate_lock_specs
 from conda_lock.src_parser.conda_common import conda_spec_to_versioned_dep
 from conda_lock.src_parser.selectors import filter_platform_selectors
 
@@ -27,10 +26,10 @@ def parse_conda_requirement(req: str) -> Tuple[str, str]:
 
 
 def _parse_environment_file_for_platform(
-    environment_file: pathlib.Path,
     content: str,
+    category: str,
     platform: str,
-) -> LockSpecification:
+) -> List[Dependency]:
     """
     Parse dependencies from a conda environment specification for an
     assumed target platform.
@@ -44,23 +43,15 @@ def _parse_environment_file_for_platform(
     """
     filtered_content = "\n".join(filter_platform_selectors(content, platform=platform))
     env_yaml_data = yaml.safe_load(filtered_content)
-
     specs = env_yaml_data["dependencies"]
-    channels: List[str] = env_yaml_data.get("channels", [])
-
-    # These extension fields are nonstandard
-    platforms: List[str] = env_yaml_data.get("platforms", [])
-    category: str = env_yaml_data.get("category") or "main"
 
     # Split out any sub spec sections from the dependencies mapping
     mapping_specs = [x for x in specs if not isinstance(x, str)]
     specs = [x for x in specs if isinstance(x, str)]
 
     dependencies: List[Dependency] = []
     for spec in specs:
-        vdep = conda_spec_to_versioned_dep(spec, category)
-        vdep.selectors.platform = [platform]
-        dependencies.append(vdep)
+        dependencies.append(conda_spec_to_versioned_dep(spec, category))
 
     for mapping_spec in mapping_specs:
         if "pip" in mapping_spec:
@@ -88,12 +79,7 @@ def _parse_environment_file_for_platform(
             # ensure pip is in target env
             dependencies.append(parse_python_requirement("pip", manager="conda"))
 
-    return LockSpecification(
-        dependencies=dependencies,
-        channels=channels,  # type: ignore
-        platforms=platforms,
-        sources=[environment_file],
-    )
+    return dependencies
 
 
 def parse_platforms_from_env_file(environment_file: pathlib.Path) -> List[str]:
@@ -127,26 +113,20 @@ def parse_environment_file(
     with environment_file.open("r") as fo:
         content = fo.read()
 
+    env_yaml_data = yaml.safe_load(content)
+    channels: List[str] = env_yaml_data.get("channels", [])
+
+    # These extension fields are nonstandard
+    category: str = env_yaml_data.get("category") or "main"
+
     # Parse with selectors for each target platform
-    spec = aggregate_lock_specs(
-        [
-            _parse_environment_file_for_platform(
-                environment_file,
-                content,
-                platform,
-            )
-            for platform in platforms
-        ]
-    )
+    dep_map = {
+        platform: _parse_environment_file_for_platform(content, category, platform)
+        for platform in platforms
+    }
 
-    # Remove platform selectors if they apply to all targets
-    for dep in spec.dependencies:
-        if dep.selectors.platform == platforms:
-            dep.selectors.platform = None
-
-    # Use the list of rendered platforms for the output spec only if
-    # there is a dependency that is not used on all platforms.
-    # This is unlike meta.yaml because environment-yaml files can contain an
-    # internal list of platforms, which should be used as long as it
-    spec.platforms = platforms
-    return spec
+    return LockSpecification(
+        dependencies=dep_map,
+        channels=channels,  # type: ignore
+        sources=[environment_file],
+    )

conda_lock/src_parser/meta_yaml.py

@@ -7,7 +7,7 @@
 
 from conda_lock.common import get_in
 from conda_lock.models.lock_spec import Dependency, LockSpecification
-from conda_lock.src_parser.aggregation import aggregate_lock_specs
+from conda_lock.src_parser.conda_common import conda_spec_to_versioned_dep
 from conda_lock.src_parser.selectors import filter_platform_selectors
 
 
@@ -94,31 +94,37 @@ def parse_meta_yaml_file(
       selectors other than platform.
     """
 
+    if not meta_yaml_file.exists():
+        raise FileNotFoundError(f"{meta_yaml_file} not found")
+
+    with meta_yaml_file.open("r") as fo:
+        t = jinja2.Template(fo.read(), undefined=UndefinedNeverFail)
+        rendered = t.render()
+        meta_yaml_data = yaml.safe_load(rendered)
+
+    channels = get_in(["extra", "channels"], meta_yaml_data, [])
+
     # parse with selectors for each target platform
-    spec = aggregate_lock_specs(
-        [
-            _parse_meta_yaml_file_for_platform(meta_yaml_file, platform)
-            for platform in platforms
-        ]
-    )
-    # remove platform selectors if they apply to all targets
-    for dep in spec.dependencies:
-        if dep.selectors.platform == platforms:
-            dep.selectors.platform = None
+    dep_map = {
+        platform: _parse_meta_yaml_file_for_platform(meta_yaml_file, platform)
+        for platform in platforms
+    }
 
-    return spec
+    return LockSpecification(
+        dependencies=dep_map,
+        channels=channels,
+        sources=[meta_yaml_file],
+    )
 
 
 def _parse_meta_yaml_file_for_platform(
     meta_yaml_file: pathlib.Path,
     platform: str,
-) -> LockSpecification:
+) -> List[Dependency]:
     """Parse a simple meta-yaml file for dependencies, assuming the target platform.
 
     * This does not support multi-output files and will ignore all lines with selectors other than platform
     """
-    if not meta_yaml_file.exists():
-        raise FileNotFoundError(f"{meta_yaml_file} not found")
 
     with meta_yaml_file.open("r") as fo:
         filtered_recipe = "\n".join(
@@ -129,17 +135,13 @@ def _parse_meta_yaml_file_for_platform(
 
         meta_yaml_data = yaml.safe_load(rendered)
 
-    channels = get_in(["extra", "channels"], meta_yaml_data, [])
     dependencies: List[Dependency] = []
 
     def add_spec(spec: str, category: str) -> None:
         if spec is None:
             return
 
-        from .conda_common import conda_spec_to_versioned_dep
-
         dep = conda_spec_to_versioned_dep(spec, category)
-        dep.selectors.platform = [platform]
         dependencies.append(dep)
 
     def add_requirements_from_recipe_or_output(yaml_data: Dict[str, Any]) -> None:
@@ -154,9 +156,4 @@ def add_requirements_from_recipe_or_output(yaml_data: Dict[str, Any]) -> None:
     for output in get_in(["outputs"], meta_yaml_data, []):
         add_requirements_from_recipe_or_output(output)
 
-    return LockSpecification(
-        dependencies=dependencies,
-        channels=channels,
-        platforms=[platform],
-        sources=[meta_yaml_file],
-    )
+    return dependencies