Merge pull request #733 from maresb/fix-pkgs-dirs-keyerror

Fix KeyError: 'pkgs_dirs'

conda_lock/conda_solver.py

@@ -9,11 +9,19 @@
 import time
 
 from contextlib import contextmanager
-from typing import Dict, Iterable, Iterator, List, MutableSequence, Optional, Sequence
+from typing import (
+    Any,
+    Dict,
+    Iterable,
+    Iterator,
+    List,
+    Literal,
+    MutableSequence,
+    Optional,
+    Sequence,
+)
 from urllib.parse import urlsplit, urlunsplit
 
-import yaml
-
 from typing_extensions import TypedDict
 
 from conda_lock.interfaces.vendored_conda import MatchSpec
@@ -239,6 +247,38 @@ def _get_repodata_record(
     return None
 
 
+def _get_pkgs_dirs(
+    *,
+    conda: PathLike,
+    platform: str,
+    method: Optional[Literal["config", "info"]] = None,
+) -> List[pathlib.Path]:
+    """Extract the package cache directories from the conda configuration."""
+    if method is None:
+        method = "config" if is_micromamba(conda) else "info"
+    if method == "config":
+        # 'package cache' was added to 'micromamba info' in v1.4.6.
+        args = [str(conda), "config", "--json", "list", "pkgs_dirs"]
+    elif method == "info":
+        args = [str(conda), "info", "--json"]
+    env = conda_env_override(platform)
+    output = subprocess.check_output(args, env=env).decode()
+    json_object_str = extract_json_object(output)
+    json_object: dict[str, Any] = json.loads(json_object_str)
+    pkgs_dirs_list: list[str]
+    if "pkgs_dirs" in json_object:
+        pkgs_dirs_list = json_object["pkgs_dirs"]
+    elif "package cache" in json_object:
+        pkgs_dirs_list = json_object["package cache"]
+    else:
+        raise ValueError(
+            f"Unable to extract pkgs_dirs from {json_object}. "
+            "Please report this issue to the conda-lock developers."
+        )
+    pkgs_dirs = [pathlib.Path(d) for d in pkgs_dirs_list]
+    return pkgs_dirs
+
+
 def _reconstruct_fetch_actions(
     conda: PathLike, platform: str, dry_run_install: DryRunInstall
 ) -> DryRunInstall:
@@ -256,34 +296,10 @@ def _reconstruct_fetch_actions(
     link_actions = {p["name"]: p for p in dry_run_install["actions"]["LINK"]}
     fetch_actions = {p["name"]: p for p in dry_run_install["actions"]["FETCH"]}
     link_only_names = set(link_actions.keys()).difference(fetch_actions.keys())
-    if is_micromamba(conda):
-        if link_only_names:
-            args = [str(conda), "config", "list", "pkgs_dirs"]
-            pkgs_dirs = [
-                pathlib.Path(d)
-                for d in yaml.safe_load(
-                    subprocess.check_output(
-                        args, env=conda_env_override(platform)
-                    ).decode()
-                )["pkgs_dirs"]
-            ]
-        else:
-            pkgs_dirs = []
+    if link_only_names:
+        pkgs_dirs = _get_pkgs_dirs(conda=conda, platform=platform)
     else:
-        if link_only_names:
-            args = [str(conda), "info", "--json"]
-            pkgs_dirs = [
-                pathlib.Path(d)
-                for d in json.loads(
-                    extract_json_object(
-                        subprocess.check_output(
-                            args, env=conda_env_override(platform)
-                        ).decode()
-                    )
-                )["pkgs_dirs"]
-            ]
-        else:
-            pkgs_dirs = []
+        pkgs_dirs = []
 
     for link_pkg_name in link_only_names:
         link_action = link_actions[link_pkg_name]

tests/test-get-pkgs-dirs/conda-info.json

@@ -0,0 +1,104 @@
+{
+    "GID": 127,
+    "UID": 1001,
+    "active_prefix": "/home/runner/miniconda3/envs/esmvaltool-fromlock",
+    "active_prefix_name": "esmvaltool-fromlock",
+    "av_data_dir": "/home/runner/miniconda3/etc/conda",
+    "av_metadata_url_base": null,
+    "channels": [
+      "https://conda.anaconda.org/conda-forge/linux-64",
+      "https://conda.anaconda.org/conda-forge/noarch"
+    ],
+    "conda_build_version": "not installed",
+    "conda_env_version": "24.9.1",
+    "conda_location": "/home/runner/miniconda3/lib/python3.12/site-packages/conda",
+    "conda_prefix": "/home/runner/miniconda3",
+    "conda_shlvl": 2,
+    "conda_version": "24.9.1",
+    "config_files": [
+      "/home/runner/miniconda3/.condarc",
+      "/home/runner/.condarc"
+    ],
+    "default_prefix": "/home/runner/miniconda3/envs/esmvaltool-fromlock",
+    "env_vars": {
+      "CIO_TEST": "<not set>",
+      "CONDA": "/home/runner/miniconda3",
+      "CONDA_DEFAULT_ENV": "esmvaltool-fromlock",
+      "CONDA_EXE": "/home/runner/miniconda3/bin/conda",
+      "CONDA_PKGS_DIR": "/home/runner/conda_pkgs_dir",
+      "CONDA_PREFIX": "/home/runner/miniconda3/envs/esmvaltool-fromlock",
+      "CONDA_PREFIX_1": "/home/runner/miniconda3",
+      "CONDA_PROMPT_MODIFIER": "",
+      "CONDA_PYTHON_EXE": "/home/runner/miniconda3/bin/python",
+      "CONDA_ROOT": "/home/runner/miniconda3",
+      "CONDA_SHLVL": "2",
+      "CURL_CA_BUNDLE": "<not set>",
+      "DEPLOYMENT_BASEPATH": "/opt/runner",
+      "GITHUB_EVENT_PATH": "/home/runner/work/_temp/_github_workflow/event.json",
+      "GITHUB_PATH": "/home/runner/work/_temp/_runner_file_commands/add_path_56b7c011-2784-49f4-8a76-8ec97d34b1bc",
+      "LD_PRELOAD": "<not set>",
+      "PATH": "/home/runner/miniconda3/envs/esmvaltool-fromlock/bin:/home/runner/miniconda3/condabin:/home/runner/miniconda3/condabin:/snap/bin:/home/runner/.local/bin:/opt/pipx_bin:/home/runner/.cargo/bin:/home/runner/.config/composer/vendor/bin:/usr/local/.ghcup/bin:/home/runner/.dotnet/tools:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/usr/games:/usr/local/games:/snap/bin:/home/runner/.dotnet/tools",
+      "REQUESTS_CA_BUNDLE": "<not set>",
+      "SELENIUM_JAR_PATH": "/usr/share/java/selenium-server.jar",
+      "SSL_CERT_FILE": "<not set>",
+      "SWIFT_PATH": "/usr/share/swift/usr/bin"
+    },
+    "envs": [
+      "/home/runner/miniconda3",
+      "/home/runner/miniconda3/envs/esmvaltool-fromlock"
+    ],
+    "envs_dirs": [
+      "/home/runner/miniconda3/envs",
+      "/home/runner/.conda/envs"
+    ],
+    "netrc_file": null,
+    "offline": false,
+    "pkgs_dirs": [
+      "/home/runner/conda_pkgs_dir"
+    ],
+    "platform": "linux-64",
+    "python_version": "3.12.6.final.0",
+    "rc_path": "/home/runner/.condarc",
+    "requests_version": "2.32.3",
+    "root_prefix": "/home/runner/miniconda3",
+    "root_writable": true,
+    "site_dirs": [],
+    "solver": {
+      "default": true,
+      "name": "libmamba",
+      "user_agent": "solver/libmamba conda-libmamba-solver/24.7.0 libmambapy/1.5.9"
+    },
+    "sys.executable": "/home/runner/miniconda3/bin/python",
+    "sys.prefix": "/home/runner/miniconda3",
+    "sys.version": "3.12.6 | packaged by conda-forge | (main, Sep 22 2024, 14:16:49) [GCC 13.3.0]",
+    "sys_rc_path": "/home/runner/miniconda3/.condarc",
+    "user_agent": "conda/24.9.1 requests/2.32.3 CPython/3.12.6 Linux/6.5.0-1025-azure ubuntu/22.04.5 glibc/2.35 solver/libmamba conda-libmamba-solver/24.7.0 libmambapy/1.5.9",
+    "user_rc_path": "/home/runner/.condarc",
+    "virtual_pkgs": [
+      [
+        "__archspec",
+        "1",
+        "zen2"
+      ],
+      [
+        "__conda",
+        "24.9.1",
+        "0"
+      ],
+      [
+        "__glibc",
+        "2.35",
+        "0"
+      ],
+      [
+        "__linux",
+        "6.5.0",
+        "0"
+      ],
+      [
+        "__unix",
+        "0",
+        "0"
+      ]
+    ]
+  }

tests/test-get-pkgs-dirs/mamba-info.json

@@ -0,0 +1,32 @@
+{
+    "base environment": "/home/runner/miniconda3/envs/esmvaltool-fromlock",
+    "channels": [
+        "https://conda.anaconda.org/conda-forge/linux-64",
+        "https://conda.anaconda.org/conda-forge/noarch"
+    ],
+    "curl version": "libcurl/8.9.1 OpenSSL/3.3.2 zlib/1.3.1 zstd/1.5.6 libssh2/1.11.0 nghttp2/1.58.0",
+    "env location": "/home/runner/miniconda3/envs/esmvaltool-fromlock",
+    "environment": "base (active)",
+    "envs directories": [
+        "/home/runner/miniconda3/envs/esmvaltool-fromlock/envs"
+    ],
+    "libarchive version": "libarchive 3.7.4 zlib/1.2.13 liblzma/5.2.6 bz2lib/1.0.8 liblz4/1.9.3 libzstd/1.5.6",
+    "libmamba version": "2.0.2",
+    "mamba version": "2.0.2",
+    "package cache": [
+        "/home/runner/conda_pkgs_dir"
+    ],
+    "platform": "linux-64",
+    "populated config files": [
+        "/home/runner/.condarc"
+    ],
+    "user config files": [
+        "/home/runner/.mambarc"
+    ],
+    "virtual packages": [
+        "__unix=0=0",
+        "__linux=6.5.0=0",
+        "__glibc=2.35=0",
+        "__archspec=1=x86_64_v3"
+    ]
+}

tests/test-get-pkgs-dirs/micromamba-1.4.5-info.json

@@ -0,0 +1,28 @@
+{
+    "base environment": "/home/user/micromamba",
+    "channels": [
+        "https://conda.anaconda.org/conda-forge/linux-64",
+        "https://conda.anaconda.org/conda-forge/noarch",
+        "https://conda.anaconda.org/nodefaults/linux-64",
+        "https://conda.anaconda.org/nodefaults/noarch"
+    ],
+    "curl version": "libcurl/7.88.1 OpenSSL/3.1.1 zlib/1.2.13 zstd/1.5.2 libssh2/1.11.0 nghttp2/1.52.0",
+    "env location": "/home/user/micromamba/envs/esmvaltool-lock-test",
+    "environment": "esmvaltool-lock-test (active)",
+    "libarchive version": "libarchive 3.6.2 zlib/1.2.13 bz2lib/1.0.8 libzstd/1.5.2",
+    "libmamba version": "1.4.5",
+    "micromamba version": "1.4.5",
+    "platform": "linux-64",
+    "populated config files": [
+        "/home/user/.condarc"
+    ],
+    "user config files": [
+        "/home/user/.mambarc"
+    ],
+    "virtual packages": [
+        "__unix=0=0",
+        "__linux=5.15.0=0",
+        "__glibc=2.31=0",
+        "__archspec=1=x86_64"
+    ]
+}

tests/test-get-pkgs-dirs/micromamba-1.4.6-info.json

@@ -0,0 +1,35 @@
+{
+    "base environment": "/home/user/micromamba",
+    "channels": [
+        "https://conda.anaconda.org/conda-forge/linux-64",
+        "https://conda.anaconda.org/conda-forge/noarch",
+        "https://conda.anaconda.org/nodefaults/linux-64",
+        "https://conda.anaconda.org/nodefaults/noarch"
+    ],
+    "curl version": "libcurl/7.88.1 OpenSSL/3.1.1 zlib/1.2.13 zstd/1.5.2 libssh2/1.11.0 nghttp2/1.52.0",
+    "env location": "/home/user/micromamba/envs/esmvaltool-lock-test",
+    "environment": "esmvaltool-lock-test (active)",
+    "envs directories": [
+        "/home/user/micromamba/envs"
+    ],
+    "libarchive version": "libarchive 3.6.2 zlib/1.2.13 bz2lib/1.0.8 libzstd/1.5.2",
+    "libmamba version": "1.4.6",
+    "micromamba version": "1.4.6",
+    "package cache": [
+        "/home/user/micromamba/pkgs",
+        "/home/user/.mamba/pkgs"
+    ],
+    "platform": "linux-64",
+    "populated config files": [
+        "/home/user/.condarc"
+    ],
+    "user config files": [
+        "/home/user/.mambarc"
+    ],
+    "virtual packages": [
+        "__unix=0=0",
+        "__linux=5.15.0=0",
+        "__glibc=2.31=0",
+        "__archspec=1=x86_64"
+    ]
+}

tests/test-get-pkgs-dirs/micromamba-config-list-pkgs_dirs.json

@@ -0,0 +1,6 @@
+{
+    "pkgs_dirs": [
+        "/home/user/micromamba/pkgs",
+        "/home/user/.mamba/pkgs"
+    ]
+}

tests/test_conda_lock.py

@@ -15,7 +15,8 @@
 
 from glob import glob
 from pathlib import Path
-from typing import Any, ContextManager, Dict, List, Literal, Set, Tuple, Union
+from typing import Any, ContextManager, Dict, List, Literal, Optional, Set, Tuple, Union
+from unittest import mock
 from unittest.mock import MagicMock
 from urllib.parse import urldefrag, urlsplit
 
@@ -47,7 +48,11 @@
     render_lockfile_for_platform,
     run_lock,
 )
-from conda_lock.conda_solver import extract_json_object, fake_conda_environment
+from conda_lock.conda_solver import (
+    _get_pkgs_dirs,
+    extract_json_object,
+    fake_conda_environment,
+)
 from conda_lock.errors import (
     ChannelAggregationError,
     MissingEnvVarError,
@@ -2815,6 +2820,49 @@ def test_extract_json_object():
     assert extract_json_object('{"key1": true }') == '{"key1": true }'
 
 
+def test_get_pkgs_dirs(conda_exe):
+    # If it runs without raising an exception, then it found the package directories.
+    _get_pkgs_dirs(conda=conda_exe, platform="linux-64")
+
+
+@pytest.mark.parametrize(
+    "info_file,expected",
+    [
+        ("conda-info.json", [Path("/home/runner/conda_pkgs_dir")]),
+        ("mamba-info.json", [Path("/home/runner/conda_pkgs_dir")]),
+        ("micromamba-1.4.5-info.json", None),
+        (
+            "micromamba-1.4.6-info.json",
+            [Path("/home/user/micromamba/pkgs"), Path("/home/user/.mamba/pkgs")],
+        ),
+        (
+            "micromamba-config-list-pkgs_dirs.json",
+            [Path("/home/user/micromamba/pkgs"), Path("/home/user/.mamba/pkgs")],
+        ),
+    ],
+)
+def test_get_pkgs_dirs_mocked_output(info_file: str, expected: Optional[List[Path]]):
+    """Test _get_pkgs_dirs with mocked subprocess.check_output."""
+    info_path = TESTS_DIR / "test-get-pkgs-dirs" / info_file
+    command_output = info_path.read_bytes()
+    conda = info_path.stem.split("-")[0]
+    method: Literal["config", "info"]
+    if "config" in info_file:
+        method = "config"
+    elif "info" in info_file:
+        method = "info"
+    else:
+        raise ValueError(f"Unknown method for {info_file}")
+
+    with mock.patch("subprocess.check_output", return_value=command_output):
+        # If expected is None, we expect a ValueError to be raised
+        with pytest.raises(
+            ValueError
+        ) if expected is None else contextlib.nullcontext():
+            result = _get_pkgs_dirs(conda=conda, platform="linux-64", method=method)
+            assert result == expected
+
+
 def test_cli_version(capsys: "pytest.CaptureFixture[str]"):
     """It should correctly report its version."""