diff --git a/conda_lock/conda_lock.py b/conda_lock/conda_lock.py index 1ba34bbe1..c31c19c10 100644 --- a/conda_lock/conda_lock.py +++ b/conda_lock/conda_lock.py @@ -542,7 +542,7 @@ def render_lockfile_for_platform( # noqa: C901 lockfile.toposort_inplace() for p in lockfile.package: - if p.platform == platform and p.category in categories: + if p.platform == platform and not p.categories.isdisjoint(categories): if p.manager == "pip": pip_deps.append(p) elif p.manager == "conda": @@ -550,9 +550,7 @@ def render_lockfile_for_platform( # noqa: C901 if not p.name.startswith("__"): conda_deps.append(p) - def format_pip_requirement( - spec: LockedDependency, platform: str, direct: bool = False - ) -> str: + def format_pip_requirement(spec: LockedDependency, direct: bool = False) -> str: if spec.source and spec.source.type == "url": return f"{spec.name} @ {spec.source.url}" elif direct: @@ -566,9 +564,7 @@ def format_pip_requirement( s += f" --hash=sha256:{spec.hash.sha256}" return s - def format_conda_requirement( - spec: LockedDependency, platform: str, direct: bool = False - ) -> str: + def format_conda_requirement(spec: LockedDependency, direct: bool = False) -> str: if direct: # inject the environment variables in here return posixpath.expandvars(f"{spec.url}#{spec.hash.md5}") @@ -589,7 +585,7 @@ def format_conda_requirement( ), "dependencies:", *( - f" - {format_conda_requirement(dep, platform, direct=False)}" + f" - {format_conda_requirement(dep, direct=False)}" for dep in conda_deps ), ] @@ -598,7 +594,7 @@ def format_conda_requirement( [ " - pip:", *( - f" - {format_pip_requirement(dep, platform, direct=False)}" + f" - {format_pip_requirement(dep, direct=False)}" for dep in pip_deps ), ] @@ -609,7 +605,7 @@ def format_conda_requirement( lockfile_contents.append("@EXPLICIT\n") lockfile_contents.extend( - [format_conda_requirement(dep, platform, direct=True) for dep in conda_deps] + [format_conda_requirement(dep, direct=True) for dep in conda_deps] ) def sanitize_lockfile_line(line: str) -> str: @@ -623,10 +619,7 @@ def sanitize_lockfile_line(line: str) -> str: # emit an explicit requirements.txt, prefixed with '# pip ' lockfile_contents.extend( - [ - f"# pip {format_pip_requirement(dep, platform, direct=True)}" - for dep in pip_deps - ] + [f"# pip {format_pip_requirement(dep, direct=True)}" for dep in pip_deps] ) if len(pip_deps) > 0: diff --git a/conda_lock/lockfile/__init__.py b/conda_lock/lockfile/__init__.py index 760bd0380..386f0cae7 100644 --- a/conda_lock/lockfile/__init__.py +++ b/conda_lock/lockfile/__init__.py @@ -1,9 +1,18 @@ -import json import pathlib from collections import defaultdict from textwrap import dedent -from typing import Any, Collection, Dict, List, Mapping, Optional, Sequence, Set, Union +from typing import ( + Any, + Collection, + DefaultDict, + Dict, + List, + Mapping, + Optional, + Set, + Union, +) import yaml @@ -38,7 +47,6 @@ def _seperator_munge_get( def apply_categories( requested: Dict[str, Dependency], planned: Mapping[str, Union[List[LockedDependency], LockedDependency]], - categories: Sequence[str] = ("main", "dev"), convert_to_pip_names: bool = False, ) -> None: """map each package onto the root request the with the highest-priority category""" @@ -57,7 +65,6 @@ def apply_categories( # walk dependency tree to assemble all transitive dependencies by request dependents: Dict[str, Set[str]] = {} - by_category = defaultdict(list) def extract_planned_items( planned_items: Union[List[LockedDependency], LockedDependency] @@ -78,8 +85,8 @@ def dep_name(manager: str, dep: str) -> str: return conda_name_to_pypi_name(dep).lower() return dep - for name, request in requested.items(): - todo: List[str] = list() + for name in requested: + todo: List[str] = [] deps: Set[str] = set() item = name @@ -106,29 +113,22 @@ def dep_name(manager: str, dep: str) -> str: dependents[name] = deps - by_category[request.category].append(request.name) - - # now, map each package to its root request preferring the ones earlier in the - # list - 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 + # now, map each package to its root requests / dependencies + root_requests: DefaultDict[str, Set[str]] = defaultdict(set) + for root, transitive_deps in dependents.items(): + for transitive_dep in transitive_deps: + root_requests[transitive_dep].add(root) + # include root requests themselves for name in requested: - root_requests[name] = name + root_requests[name].add(name) - for dep, root in root_requests.items(): - source = requested[root] - # 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 dep, roots in root_requests.items(): + target = _seperator_munge_get(planned, dep) + for root in roots: + source = requested[root] + assert isinstance(target, LockedDependency) # TODO: why? + target.categories.add(source.category) def parse_conda_lock_file(path: pathlib.Path) -> Lockfile: @@ -141,10 +141,13 @@ def parse_conda_lock_file(path: pathlib.Path) -> Lockfile: if not (isinstance(version, int) and version <= Lockfile.version): raise ValueError(f"{path} has unknown version {version}") + packages = {} for p in content["package"]: + del p["category"] del p["optional"] + packages[(p["name"], p["version"], p["platform"])] = p - return Lockfile.parse_obj(content) + return Lockfile.parse_obj({**content, "package": list(packages.values())}) def write_conda_lock_file( @@ -156,7 +159,7 @@ def write_conda_lock_file( content.toposort_inplace() with path.open("w") as f: if include_help_text: - categories = set(p.category for p in content.package) + categories = {cat for p in content.package for cat in p.categories} def write_section(text: str) -> None: lines = dedent(text).split("\n") @@ -209,20 +212,28 @@ def write_section(text: str) -> None: output: Dict[str, Any] = { "version": Lockfile.version, - "metadata": json.loads( - content.metadata.json( - by_alias=True, exclude_unset=True, exclude_none=True - ) + "metadata": content.metadata.dict( + by_alias=True, exclude_unset=True, exclude_none=True ), - "package": [ - { - **package.dict( - by_alias=True, exclude_unset=True, exclude_none=True - ), - "optional": (package.category != "main"), - } - for package in content.package - ], + "package": [], } + for package in content.package: + sorted_cats = sorted(package.categories) + for category in sorted_cats: + output["package"].append( + dict( + sorted( + { + **package.dict( + by_alias=True, exclude_unset=True, exclude_none=True + ), + "categories": sorted_cats, + "category": category, + "optional": (category != "main"), + }.items() + ) + ) + ) + yaml.dump(output, stream=f, sort_keys=False) diff --git a/conda_lock/lockfile/models.py b/conda_lock/lockfile/models.py index 296a98023..0982dd579 100644 --- a/conda_lock/lockfile/models.py +++ b/conda_lock/lockfile/models.py @@ -6,7 +6,16 @@ import typing from collections import defaultdict, namedtuple -from typing import TYPE_CHECKING, AbstractSet, ClassVar, Dict, List, Optional, Union +from typing import ( + TYPE_CHECKING, + AbstractSet, + ClassVar, + Dict, + List, + Optional, + Set, + Union, +) if TYPE_CHECKING: @@ -44,7 +53,7 @@ class LockedDependency(StrictModel): dependencies: Dict[str, str] = {} url: str hash: HashModel - category: str = "main" + categories: Set[str] = set() source: Optional[DependencySource] = None build: Optional[str] = None diff --git a/tests/gdal/environment.yml b/tests/gdal/environment.yml index 58efbe4bf..dbae97604 100644 --- a/tests/gdal/environment.yml +++ b/tests/gdal/environment.yml @@ -6,4 +6,4 @@ dependencies: - python >= 3.7, < 3.8 - gdal - pip: - - toolz \ No newline at end of file + - toolz diff --git a/tests/test_conda_lock.py b/tests/test_conda_lock.py index ba82b43cf..d4aa3d5b4 100644 --- a/tests/test_conda_lock.py +++ b/tests/test_conda_lock.py @@ -33,6 +33,7 @@ _add_auth_to_line, _add_auth_to_lockfile, _extract_domain, + _solve_for_arch, _strip_auth_from_line, _strip_auth_from_lockfile, create_lockfile_from_spec, @@ -346,7 +347,7 @@ def test_lock_poetry_ibis( all_categories = set() for pkg in lockfile.package: - all_categories.add(pkg.category) + all_categories.update(pkg.categories) for desired_category in extra_categories: assert ( @@ -520,6 +521,7 @@ def test_choose_wheel() -> None: platform="linux-64", ) assert len(solution) == 1 + assert solution["fastavro"].categories == {"main"} assert solution["fastavro"].hash == HashModel( sha256="a111a384a786b7f1fd6a8a8307da07ccf4d4c425084e2d61bae33ecfb60de405" ) @@ -1183,6 +1185,54 @@ def test_run_lock_with_input_hash_check( assert "Spec hash already locked for" in output.err +def test_solve_arch_multiple_categories(): + _conda_exe = determine_conda_executable(None, mamba=False, micromamba=False) + vpr = default_virtual_package_repodata() + channels = [Channel.from_string("conda-forge")] + + with vpr, tempfile.NamedTemporaryFile(dir=".") as tf: + spec = LockSpecification( + dependencies={ + "linux-64": [ + VersionedDependency( + name="python", + version="=3.10.9", + manager="conda", + category="main", + extras=[], + ), + VersionedDependency( + name="pandas", + version="=1.5.3", + manager="conda", + category="test", + extras=[], + ), + VersionedDependency( + name="pyarrow", + version="=9.0.0", + manager="conda", + category="dev", + extras=[], + ), + ], + }, + channels=channels, + # NB: this file must exist for relative path resolution to work + # in create_lockfile_from_spec + sources=[Path(tf.name)], + virtual_package_repo=vpr, + ) + + locked_deps = _solve_for_arch(_conda_exe, spec, "linux-64", channels) + python_deps = [dep for dep in locked_deps if dep.name == "python"] + assert len(python_deps) == 1 + assert python_deps[0].categories == {"main", "test", "dev"} + numpy_deps = [dep for dep in locked_deps if dep.name == "numpy"] + assert len(numpy_deps) == 1 + assert numpy_deps[0].categories == {"test", "dev"} + + @pytest.mark.parametrize( "package,version,url_pattern", [