qclib remove CLI command now supports --only-ref option

Author: aabmets

quantcrypt/internal/cli/annotations.py

@@ -19,6 +19,7 @@
     "CompileAlgos",
     "RemoveAlgos",
     "KeepAlgos",
+    "OnlyRef",
     "WithOpt",
     "Version",
     "DryRun",
@@ -74,6 +75,14 @@
     ])
 )]
 
+OnlyRef = Annotated[bool, Option(
+    "--only-ref", "-r", show_default=False, help=' '.join([
+        "Can be used together with the --keep option to keep only the clean reference binaries",
+        "of algorithms. Useful for when QuantCrypt is being used within virtualized environments",
+        "which do not support specialized CPU instructions."
+    ])
+)]
+
 WithOpt = Annotated[bool, Option(
     "--with-opt", "-o", show_default=False, help=' '.join([
         "Includes architecture-specific optimized variants to compilation targets.",

quantcrypt/internal/cli/commands/remove.py

@@ -9,6 +9,7 @@
 #   SPDX-License-Identifier: MIT
 #
 
+import json
 from typer import Typer
 from itertools import product
 from quantcrypt.internal import utils, constants as const
@@ -24,43 +25,101 @@
 )
 
 
+def remove_spec_variants(
+        spec_variants: dict[const.AlgoSpec, list[const.PQAVariant]]
+) -> tuple[dict, dict]:
+    removed_variants: dict[const.AlgoSpec, list[const.PQAVariant]] = dict()
+    already_removed: dict[const.AlgoSpec, list[const.PQAVariant]] = dict()
+    bin_path = utils.search_upwards("bin")
+    bin_contents = list(bin_path.iterdir())
+
+    for spec, variants in spec_variants.items():  # type: const.AlgoSpec, list[const.PQAVariant]
+        for variant in variants:
+            did_remove = False
+
+            for item in bin_contents:
+                if spec.module_name(variant) in item.name and item.exists():
+                    item.unlink()
+                    x = removed_variants.get(spec, list())
+                    x.append(variant)
+                    removed_variants[spec] = x
+                    did_remove = True
+
+            if not did_remove:
+                y = already_removed.get(spec, list())
+                y.append(variant)
+                already_removed[spec] = y
+
+    return removed_variants, already_removed
+
+
+def report_spec_variants(
+        spec_variants: dict[const.AlgoSpec, list[const.PQAVariant]]
+) -> None:
+    armor_names = [s.armor_name() for s in spec_variants.keys()]
+    longest_name_len = max(len(n) for n in armor_names) if armor_names else 0
+
+    for spec, variants in spec_variants.items():
+        variants_fmt = json.dumps([v.value for v in variants])
+        arna_fmt = spec.armor_name().rjust(longest_name_len)
+        console.styled_print(f"{arna_fmt}: {variants_fmt}")
+
+
 @remove_app.callback()
 def command_remove(
         algorithms: ats.RemoveAlgos,
         keep_algos: ats.KeepAlgos = False,
+        only_ref: ats.OnlyRef = False,
         dry_run: ats.DryRun = False,
         non_interactive: ats.NonInteractive = False
 ) -> None:
+    if only_ref and not keep_algos:
+        console.raise_error("Cannot use --only-ref without --keep")
+
+    chosen_algos = const.SupportedAlgos.filter(algorithms)
+    if len(chosen_algos) != len(algorithms):
+        algo_names = [s.armor_name() for s in chosen_algos]
+        bad_names = [a for a in algorithms if a.upper() not in algo_names]
+        if bad_names:  # pragma: no branch
+            console.raise_error(
+                f"Unknown algorithm name(s): {json.dumps(bad_names)}. " +
+                "Please choose algorithm names from the following list:\n" +
+                ' | '.join(const.SupportedAlgos.armor_names())
+            )
+
     console.notify_dry_run(dry_run)
     console.styled_print("QuantCrypt is about to remove compiled PQA binaries from itself.")
 
     if not non_interactive:
         console.ask_continue(exit_on_false=True)
 
-    algos = const.SupportedAlgos.filter(algorithms, invert=keep_algos)
     variants = const.PQAVariant.members()
+    if only_ref:
+        algorithms = [a.upper() for a in algorithms]
+        algos = const.SupportedAlgos
+    else:
+        algos = const.SupportedAlgos.filter(algorithms, invert=keep_algos)
+
+    to_remove: dict[const.AlgoSpec, list[const.PQAVariant]] = dict()
+    for spec, variant in product(algos, variants):  # type: const.AlgoSpec, const.PQAVariant
+        if only_ref and spec.armor_name() in algorithms and variant == const.PQAVariant.REF:
+            continue
+        variants = to_remove.get(spec, list())
+        variants.append(variant)
+        to_remove[spec] = variants
 
     if dry_run:
-        console.styled_print("QuantCrypt would have removed the following algorithms:")
-        console.pretty_print(', '.join(s.armor_name() for s in algos))
+        console.styled_print("\nQuantCrypt would have removed the following algorithms and their variants:")
+        report_spec_variants(to_remove)
         return
 
-    print()
-    bin_path = utils.search_upwards("bin")
-    bin_contents = list(bin_path.iterdir())
-
-    for spec, variant in product(algos, variants):  # type: const.AlgoSpec, const.PQAVariant
-        module_name = spec.module_name(variant)
-        did_remove = False
+    removed_variants, already_removed = remove_spec_variants(to_remove)
 
-        for item in bin_contents:
-            if spec.module_name(variant) in item.name and item.exists():
-                item.unlink()
-                did_remove = True
-                console.styled_print(f"Successfully removed binaries - {module_name}")
+    console.styled_print("\nSuccessfully removed binaries: ")
+    report_spec_variants(removed_variants)
 
-        if not did_remove:
-            console.styled_print(f"Binaries already removed - {module_name}")
+    console.styled_print("\nAlready removed binaries: ")
+    report_spec_variants(already_removed)
 
     print()
     console.print_success()

quantcrypt/internal/constants.py

@@ -133,9 +133,10 @@ def armor_names(self: list[AlgoSpec], pqa_type: PQAType | None = None) -> list[s
         ]
 
     def filter(self, armor_names: list[str], invert: bool = False) -> list[AlgoSpec]:
+        armor_names = [n.upper() for n in armor_names]
         return [
             spec for spec, name in product(self, armor_names)
-            if (not invert and spec.armor_name() == name.upper())
+            if (not invert and spec.armor_name() == name)
             or (invert and spec.armor_name() not in armor_names)
         ]
 

tests/test_cli/test_remove.py

@@ -9,6 +9,7 @@
 #   SPDX-License-Identifier: MIT
 #
 
+from pathlib import Path
 from unittest.mock import patch
 from quantcrypt.internal import constants as const
 from .conftest import CLIMessages
@@ -18,11 +19,21 @@ def test_remove(cli_runner, alt_tmp_path) -> None:
     with patch('internal.cli.commands.remove.utils.search_upwards') as mock:
         mock.return_value = alt_tmp_path
 
+        cli_runner("remove", ["mlkem512"], "n\n", CLIMessages.CANCELLED)
+        cli_runner("remove", ["-D", "mlkem512"], "y\n", CLIMessages.DRYRUN)
+        cli_runner("remove", ["--only-ref", "mlkem512"], "", CLIMessages.ERROR)
+
         for spec in const.SupportedAlgos:  # type: const.AlgoSpec
-            (alt_tmp_path / spec.module_name(const.PQAVariant.REF)).touch()
+            (alt_tmp_path / spec.module_name(const.PQAVariant.REF)).touch(exist_ok=True)
 
-        cli_runner("remove", ["mlkem512"], "n\n", CLIMessages.CANCELLED)
         cli_runner("remove", ["mlkem512"], "y\n", CLIMessages.SUCCESS)
-        cli_runner("remove", ["asdfg"], "y\n", CLIMessages.SUCCESS)
-        cli_runner("remove", ["-D", "mlkem512"], "y\n", CLIMessages.DRYRUN)
+
+        spec = const.SupportedAlgos.filter(["mlkem512"])[0]
+        for item in alt_tmp_path.iterdir():  # type: Path
+            for variant in const.PQAVariant.members():
+                assert spec.module_name(variant) not in item.name
+
+        cli_runner("remove", ["asdfg"], "y\n", CLIMessages.ERROR)
         cli_runner("remove", ["-N", "fastsphincs"], "", CLIMessages.SUCCESS)
+        cli_runner("remove", ["--keep", "--only-ref", "mlkem512"], "y\n", CLIMessages.SUCCESS)
+