add subcommand plugin for repoquery (#258)

* add subcommand plugin for repoquery * pre-commit-pr * add test * add news * skip test_reorder_channel_priority for now * pre-commit * CLI UX improvements * add docs * pre-commit * Apply suggestions from code review Co-authored-by: Jannis Leidel <jannis@leidel.info> --------- Co-authored-by: Jannis Leidel <jannis@leidel.info>
conda · Aug 29, 2023 · 1443a46 · 1443a46
1 parent 957b221
commit 1443a46
Show file tree

Hide file tree

Showing 8 changed files with 244 additions and 8 deletions.
diff --git a/conda_libmamba_solver/index.py b/conda_libmamba_solver/index.py
@@ -111,6 +111,7 @@ def __init__(
         channels: Iterable[Union[Channel, str]] = None,
         subdirs: Iterable[str] = None,
         repodata_fn: str = REPODATA_FN,
+        query_format=api.QueryFormat.JSON,
     ):
         self._channels = context.channels if channels is None else channels
         self._subdirs = context.subdirs if subdirs is None else subdirs
@@ -126,7 +127,7 @@ def __init__(
         self._repos += [info.repo for info in self._index.values()]
 
         self._query = api.Query(self._pool)
-        self._format = api.QueryFormat.JSON
+        self._format = query_format
 
     def get_info(self, key: str) -> _ChannelRepoInfo:
         orig_key = key
@@ -291,17 +292,23 @@ def _load_installed(self, records: Iterable[PackageRecord]) -> api.Repo:
         repo.set_installed()
         return repo
 
-    def whoneeds(self, query: str, records=True) -> Union[Iterable[PackageRecord], dict]:
+    def whoneeds(self, query: str, records=True) -> Union[Iterable[PackageRecord], dict, str]:
         result_str = self._query.whoneeds(query, self._format)
-        return self._process_query_result(result_str, records=records)
+        if self._format == api.QueryFormat.JSON:
+            return self._process_query_result(result_str, records=records)
+        return result_str
 
-    def depends(self, query: str, records=True) -> Union[Iterable[PackageRecord], dict]:
+    def depends(self, query: str, records=True) -> Union[Iterable[PackageRecord], dict, str]:
         result_str = self._query.depends(query, self._format)
-        return self._process_query_result(result_str, records=records)
+        if self._format == api.QueryFormat.JSON:
+            return self._process_query_result(result_str, records=records)
+        return result_str
 
-    def search(self, query: str, records=True) -> Union[Iterable[PackageRecord], dict]:
+    def search(self, query: str, records=True) -> Union[Iterable[PackageRecord], dict, str]:
         result_str = self._query.find(query, self._format)
-        return self._process_query_result(result_str, records=records)
+        if self._format == api.QueryFormat.JSON:
+            return self._process_query_result(result_str, records=records)
+        return result_str
 
     def explicit_pool(self, specs: Iterable[MatchSpec]) -> Iterable[str]:
         """

diff --git a/conda_libmamba_solver/plugin.py b/conda_libmamba_solver/plugin.py
@@ -3,6 +3,7 @@
 # SPDX-License-Identifier: BSD-3-Clause
 from conda import plugins
 
+from .repoquery import configure_parser, repoquery
 from .solver import LibMambaSolver
 
 
@@ -15,3 +16,13 @@ def conda_solvers():
         name="libmamba",
         backend=LibMambaSolver,
     )
+
+
+@plugins.hookimpl
+def conda_subcommands():
+    yield plugins.CondaSubcommand(
+        name="repoquery",
+        summary="Advanced search for repodata.",
+        action=repoquery,
+        configure_parser=configure_parser,
+    )
diff --git a/conda_libmamba_solver/repoquery.py b/conda_libmamba_solver/repoquery.py
@@ -0,0 +1,159 @@
+# Copyright (C) 2019, QuantStack
+# Copyright (C) 2022 Anaconda, Inc
+# Copyright (C) 2023 conda
+# SPDX-License-Identifier: BSD-3-Clause
+import argparse
+import json
+import sys
+
+from conda.base.context import context
+from conda.cli import conda_argparse
+from conda.common.io import Spinner
+from conda.core.prefix_data import PrefixData
+from libmambapy import QueryFormat
+
+from .index import LibMambaIndexHelper
+
+
+def configure_parser(parser: argparse.ArgumentParser):
+    package_cmds = argparse.ArgumentParser(add_help=False)
+    package_cmds.add_argument("package_query", help="The target package.")
+    package_grp = package_cmds.add_argument_group("Subcommand options")
+    package_grp.add_argument(
+        "-i",
+        "--installed",
+        action="store_true",
+        default=True,
+        help=argparse.SUPPRESS,
+    )
+
+    package_grp.add_argument(
+        "-p",
+        "--platform",
+        default=context.subdir,
+        help="Platform/subdir to search packages for. Defaults to current platform.",
+    )
+    package_grp.add_argument(
+        "--no-installed", action="store_true", help="Do not search currently installed packages."
+    )
+    package_grp.add_argument(
+        "--pretty", action="store_true", help="Prettier output with more details."
+    )
+
+    package_grp.add_argument(
+        "-a",
+        "--all-channels",
+        action="store_true",
+        help="Look at all channels (for depends / whoneeds).",
+    )
+
+    view_cmds = argparse.ArgumentParser(add_help=False)
+    view_grp = view_cmds.add_argument_group("Dependency options")
+    view_grp.add_argument(
+        "-t", "--tree", action="store_true", help="Show dependencies in a tree-like format."
+    )
+    view_grp.add_argument("--recursive", action="store_true", help="Show dependencies recursively.")
+
+    subparser = parser.add_subparsers(dest="subcmd")
+
+    whoneeds = subparser.add_parser(
+        "whoneeds",
+        help="Show packages that depend on this package.",
+        parents=[package_cmds, view_cmds],
+    )
+
+    depends = subparser.add_parser(
+        "depends",
+        help="Show dependencies of this package.",
+        parents=[package_cmds, view_cmds],
+    )
+
+    search = subparser.add_parser(
+        "search",
+        help="Show all available package versions.",
+        parents=[package_cmds],
+    )
+
+    for cmd in (whoneeds, search, depends):
+        conda_argparse.add_parser_channels(cmd)
+        conda_argparse.add_parser_networking(cmd)
+        conda_argparse.add_parser_known(cmd)
+        conda_argparse.add_parser_json(cmd)
+
+
+def repoquery(args):
+    if not args.subcmd:
+        print("repoquery needs a subcommand (search, depends or whoneeds), e.g.:", file=sys.stderr)
+        print("    conda repoquery search python\n", file=sys.stderr)
+        return 1
+
+    cli_flags = [getattr(args, attr, False) for attr in ("tree", "recursive", "pretty")]
+    if sum([context.json, *cli_flags]) > 1:
+        print("Use only one of --json, --tree, --recursive and --pretty.", file=sys.stderr)
+        return 1
+
+    if hasattr(args, "channel"):
+        channels = args.channel
+    else:
+        channels = None
+    if args.all_channels or (channels is None and args.subcmd == "search"):
+        if channels:
+            print("WARNING: Using all channels instead of configured channels\n", file=sys.stderr)
+        channels = context.channels
+
+    use_installed = args.installed
+    if args.no_installed:
+        use_installed = False
+
+    # if we're asking for depends and channels are given, disregard
+    # installed packages to prevent weird mixing
+    if args.subcmd in ("depends", "whoneeds") and use_installed and channels:
+        use_installed = False
+
+    if args.subcmd == "search" and not args.installed:
+        only_installed = False
+    elif args.all_channels or (channels and len(channels)):
+        only_installed = False
+    else:
+        only_installed = True
+
+    if only_installed and args.no_installed:
+        print("No channels selected. Use -a to search all channels.", file=sys.stderr)
+        return 1
+
+    if use_installed:
+        prefix_data = PrefixData(context.target_prefix)
+        prefix_data.load()
+        installed_records = prefix_data.iter_records()
+    else:
+        installed_records = ()
+
+    if context.json:
+        query_format = QueryFormat.JSON
+    elif getattr(args, "tree", None):
+        query_format = QueryFormat.TREE
+    elif getattr(args, "recursive", None):
+        query_format = QueryFormat.RECURSIVETABLE
+    elif getattr(args, "pretty", None):
+        query_format = QueryFormat.PRETTY
+    else:
+        query_format = QueryFormat.TABLE
+
+    with Spinner(
+        "Collecting package metadata",
+        enabled=not context.verbosity and not context.quiet,
+        json=context.json,
+    ):
+        index = LibMambaIndexHelper(
+            installed_records=installed_records,
+            channels=channels,
+            subdirs=(args.platform, "noarch"),
+            repodata_fn=context.repodata_fns[-1],
+            query_format=query_format,
+        )
+
+    result = getattr(index, args.subcmd)(args.package_query, records=False)
+    if context.json:
+        print(json.dumps(result, indent=2))
+    else:
+        print(result)
diff --git a/dev/collect_upstream_conda_tests/collect_upstream_conda_tests.py b/dev/collect_upstream_conda_tests/collect_upstream_conda_tests.py
@@ -83,7 +83,12 @@
         "test_get_solver_backend_multiple",
     ],
     # TODO: Investigate these, since they are solver related-ish
-    "tests/conda_env/specs/test_requirements.py": ["TestRequiremets::test_environment"],
+    "tests/conda_env/specs/test_requirements.py": [
+        "TestRequirements::test_environment",
+    ],
+    # TODO: Known to fail; should be fixed by
+    # https://github.com/conda/conda-libmamba-solver/pull/242
+    "tests/test_priority.py": ["test_reorder_channel_priority"],
 }
 
 _broken_by_libmamba_1_4_2 = {

diff --git a/docs/index.md b/docs/index.md
@@ -10,6 +10,7 @@ Start here if you want a faster conda:
 ```{toctree}
 :maxdepth: 1
 getting-started
+subcommands
 libmamba-vs-classic
 performance
 faq

diff --git a/docs/subcommands.md b/docs/subcommands.md
@@ -0,0 +1,14 @@
+# Subcommands
+
+The conda-libmamba-solver package also provides conda subcommand plugins in addition to the solver plugin.
+
+## `conda repoquery`
+
+A conda subcommand plugin that offers the same functionality as `mamba repoquery` and
+`micromamba repoquery`. It provides three actions:
+
+- `conda repoquery search`: Query repodata for packages matching a pattern.
+- `conda repoquery depends`: Show the dependencies of the requested package.
+- `conda repoquery whoneeds`: Show the packages that depend on the requested package.
+
+Check the `--help` messages for each task for more information.
diff --git a/news/258-repoquery b/news/258-repoquery
@@ -0,0 +1,19 @@
+### Enhancements
+
+* Expose libmamba's `repoquery` search features as a conda subcommand plugin. (#258)
+
+### Bug fixes
+
+* <news item>
+
+### Deprecations
+
+* <news item>
+
+### Docs
+
+* <news item>
+
+### Other
+
+* <news item>
diff --git a/tests/test_repoquery.py b/tests/test_repoquery.py
@@ -0,0 +1,20 @@
+# Copyright (C) 2022 Anaconda, Inc
+# Copyright (C) 2023 conda
+# SPDX-License-Identifier: BSD-3-Clause
+import json
+
+from .utils import conda_subprocess
+
+
+def test_repoquery():
+    p = conda_subprocess("repoquery", "--help")
+    assert "whoneeds" in p.stdout
+    assert "depends" in p.stdout
+    assert "search" in p.stdout
+
+    p = conda_subprocess("repoquery", "depends", "conda", "--json")
+    print(p.stdout)
+    data = json.loads(p.stdout)
+    assert data["result"]["status"] == "OK"
+    assert len(data["result"]["pkgs"]) > 0
+    assert len([p for p in data["result"]["pkgs"] if p["name"] == "python"]) == 1