gh-300: Add uninstall cleanup mode

src/manage/commands.py

@@ -207,6 +207,7 @@ def execute(self):
     },
 
     "uninstall": {
+        "cleanup": ("cleanup", True),
         "purge": ("purge", True),
         "by-id": ("by_id", True),
         # Undocumented aliases so that install and uninstall can be mirrored
@@ -877,15 +878,19 @@ def execute(self):
 class UninstallCommand(BaseCommand):
     CMD = "uninstall"
     HELP_LINE = ("Remove one or more runtimes from your machine. Pass " +
-                 "!B!--purge!W! to clean up all runtimes and cached files.")
+                 "!B!--cleanup!W! to remove cached and unrecognized files while " +
+                 "preserving recognized runtimes, or !B!--purge!W! to remove all.")
     USAGE_LINE = "uninstall !B!<TAG>!W!"
     HELP_TEXT = r"""!G!Uninstall command!W!
 Removes one or more runtimes from your machine.
 
 > py uninstall !B![options] <TAG> [<TAG>] ...!W!
 
 !G!Options:!W!
-    --purge         Remove all runtimes, shortcuts, and cached files. Ignores tags.
+    --cleanup       Remove cached and unrecognized content while preserving
+                    recognized runtimes. Does not accept tags.
+    --purge         Remove all runtimes, shortcuts, and cached files. Does not
+                    accept tags.
     --by-id         Require TAG to exactly match the install ID. (For advanced use.)
     !B!<TAG> <TAG>!W! ... One or more runtimes to uninstall (Company\Tag format)
                     Each tag will only remove a single runtime, even if it matches
@@ -897,11 +902,15 @@ class UninstallCommand(BaseCommand):
 !B!EXAMPLE:!W! Uninstall all runtimes without confirmation
 > py uninstall --yes --purge
 
+!B!EXAMPLE:!W! Clean up cached and unrecognized content
+> py uninstall --cleanup
+
 !B!EXAMPLE:!W! Uninstall all runtimes using their install ID.
 > py uninstall --by-id (py list --only-managed -f=id)
 """
 
     confirm = True
+    cleanup = False
     purge = False
     by_id = False
 

src/manage/uninstall_command.py

@@ -17,6 +17,35 @@ def _iterdir(p, only_files=False):
         return []
 
 
+def _is_relative_to(path, root):
+    path = Path(path)
+    root = Path(root)
+    try:
+        relative = path.relative_to(root)
+    except ValueError:
+        return False
+    return root / relative == path
+
+
+def _remove_unrecognized(root, preserve, warn_msg):
+    to_remove = []
+    for path in _iterdir(root):
+        if path in preserve:
+            continue
+        nested_preserve = {p for p in preserve if _is_relative_to(p, path)}
+        if nested_preserve:
+            _remove_unrecognized(path, nested_preserve, warn_msg)
+        else:
+            to_remove.append(path)
+
+    if to_remove:
+        LOGGER.info("Removing cached and unrecognized files from %s", root)
+        for path in to_remove:
+            rmtree(path, after_5s_warning=warn_msg.format(
+                "cached and unrecognized files"
+            ))
+
+
 def _do_purge_global_dir(global_dir, warn_msg, *, hive=None, subkey="Environment"):
     import winreg
 
@@ -76,12 +105,45 @@ def execute(cmd):
         "Ensure no Python interpreters are running, and continue to wait " +
         "or press Ctrl+C to abort.")
 
+    if cmd.purge and cmd.cleanup:
+        raise ArgumentError("--purge and --cleanup cannot be combined.")
+    if (cmd.purge or cmd.cleanup) and cmd.args:
+        option = "--purge" if cmd.purge else "--cleanup"
+        raise ArgumentError(f"{option} does not accept runtime tags.")
+
     # Clear any active venv so we don't try to delete it
     cmd.virtual_env = None
     installed = list(cmd.get_installs())
 
     cmd.tags = []
 
+    if cmd.cleanup:
+        prompt = ("Clean up cached files and unrecognized content? "
+                  "This will preserve recognized runtimes.")
+        if not cmd.ask_yn(prompt):
+            LOGGER.debug("END uninstall_command.execute")
+            return
+
+        recognized = {Path(i["prefix"]) for i in installed}
+        roots = {Path(cmd.install_dir), Path(cmd.download_dir)}
+        for root in roots:
+            if any(_is_relative_to(root, other)
+                   for other in roots if other != root):
+                continue
+            if any(_is_relative_to(path, root) for path in recognized):
+                _remove_unrecognized(root, recognized, warn_msg)
+            else:
+                LOGGER.info(
+                    "Removing cached and unrecognized files from %s", root
+                )
+                rmtree(root, after_5s_warning=warn_msg.format(
+                    "cached and unrecognized files"
+                ))
+
+        update_all_shortcuts(cmd)
+        LOGGER.debug("END uninstall_command.execute")
+        return
+
     if cmd.purge:
         if not cmd.ask_yn("Uninstall all runtimes?"):
             LOGGER.debug("END uninstall_command.execute")

tests/test_commands.py

@@ -75,6 +75,21 @@ def test_legacy_list_command_args():
         commands.ListLegacyCommand(["--list", "--help"])
 
 
+def test_uninstall_cleanup_arg(assert_log):
+    cmd = commands.UninstallCommand(
+        ["uninstall", "--cleanup", "--yes"]
+    )
+
+    assert cmd.cleanup
+    assert not cmd.purge
+    assert not cmd.confirm
+    assert not cmd.args
+
+    prompt = "This prompt should be bypassed"
+    assert cmd.ask_yn(prompt)
+    assert_log(assert_log.not_logged(prompt))
+
+
 def test_legacy_listpaths_command_help(assert_log, patched_installs):
     cmd = commands.ListPathsLegacyCommand(["--list-paths"])
     cmd.help()

tests/test_uninstall_command.py

@@ -5,7 +5,7 @@
 from pathlib import Path
 
 from manage import uninstall_command as UC
-from manage.exceptions import FilesInUseError
+from manage.exceptions import ArgumentError, FilesInUseError
 
 
 def test_purge_global_dir(monkeypatch, registry, tmp_path):
@@ -22,17 +22,19 @@ def test_purge_global_dir(monkeypatch, registry, tmp_path):
 
 def test_null_purge(fake_config):
     cmd = fake_config
-    cmd.args = ["--purge"]
+    cmd.args = []
     cmd.confirm = False
     cmd.purge = True
+    cmd.cleanup = False
     UC.execute(cmd)
 
 
 def test_purge_unknown_files(fake_config):
     cmd = fake_config
-    cmd.args = ["--purge"]
+    cmd.args = []
     cmd.confirm = False
     cmd.purge = True
+    cmd.cleanup = False
 
     unknown_file = cmd.install_dir / "unknown.txt"
     unknown_file.write_bytes(b"unknown")
@@ -48,9 +50,10 @@ def test_purge_unknown_files(fake_config):
 
 def test_purge_preserves_runtime_in_use(fake_config, monkeypatch):
     cmd = fake_config
-    cmd.args = ["--purge"]
+    cmd.args = []
     cmd.confirm = False
     cmd.purge = True
+    cmd.cleanup = False
 
     runtime = cmd.install_dir / "runtime"
     runtime.mkdir()
@@ -78,3 +81,131 @@ def locked_rmtree(path, *args, **kwargs):
 
     assert executable.is_file()
     assert not unknown_dir.exists()
+
+
+def test_cleanup_preserves_recognized_runtimes(fake_config, monkeypatch):
+    cmd = fake_config
+    cmd.args = []
+    cmd.purge = False
+    cmd.cleanup = True
+
+    runtime = cmd.install_dir / "runtime"
+    runtime.mkdir()
+    executable = runtime / "python.exe"
+    executable.write_bytes(b"preserve")
+    cmd.installs = [{
+        "display-name": "Recognized runtime",
+        "prefix": runtime,
+    }]
+
+    cmd.download_dir.mkdir()
+    cached_file = cmd.download_dir / "cached.zip"
+    cached_file.write_bytes(b"cached")
+    unknown_file = cmd.install_dir / "unknown.txt"
+    unknown_file.write_bytes(b"unknown")
+    broken_runtime = cmd.install_dir / "broken-runtime"
+    broken_runtime.mkdir()
+    (broken_runtime / "__install__.json").write_text("invalid")
+
+    prompts = []
+    cmd.ask_yn = lambda prompt: prompts.append(prompt) or True
+    refreshed = []
+    monkeypatch.setattr(UC, "update_all_shortcuts", refreshed.append)
+
+    UC.execute(cmd)
+
+    assert prompts == [
+        "Clean up cached files and unrecognized content? "
+        "This will preserve recognized runtimes."
+    ]
+    assert executable.read_bytes() == b"preserve"
+    assert not cmd.download_dir.exists()
+    assert not unknown_file.exists()
+    assert not broken_runtime.exists()
+    assert refreshed == [cmd]
+
+
+def test_cleanup_shared_download_dir_preserves_recognized_runtimes(
+    fake_config, monkeypatch
+):
+    cmd = fake_config
+    cmd.args = []
+    cmd.purge = False
+    cmd.cleanup = True
+    cmd.download_dir = cmd.install_dir
+
+    runtime = cmd.install_dir / "runtime"
+    runtime.mkdir()
+    executable = runtime / "python.exe"
+    executable.write_bytes(b"preserve")
+    cmd.installs = [{
+        "display-name": "Recognized runtime",
+        "prefix": runtime,
+    }]
+
+    cached_file = cmd.download_dir / "cached.zip"
+    cached_file.write_bytes(b"cached")
+    unknown_dir = cmd.install_dir / "unknown"
+    unknown_dir.mkdir()
+
+    cmd.ask_yn = lambda prompt: True
+    refreshed = []
+    monkeypatch.setattr(UC, "update_all_shortcuts", refreshed.append)
+
+    UC.execute(cmd)
+
+    assert executable.read_bytes() == b"preserve"
+    assert not cached_file.exists()
+    assert not unknown_dir.exists()
+    assert refreshed == [cmd]
+
+
+def test_cleanup_declined_preserves_all_files(fake_config, monkeypatch):
+    cmd = fake_config
+    cmd.args = []
+    cmd.purge = False
+    cmd.cleanup = True
+
+    cmd.download_dir.mkdir()
+    cached_file = cmd.download_dir / "cached.zip"
+    cached_file.write_bytes(b"cached")
+    unknown_file = cmd.install_dir / "unknown.txt"
+    unknown_file.write_bytes(b"unknown")
+
+    prompts = []
+    cmd.ask_yn = lambda prompt: prompts.append(prompt) or False
+    monkeypatch.setattr(
+        UC,
+        "update_all_shortcuts",
+        lambda cmd: pytest.fail("shortcuts refreshed after cleanup was declined"),
+    )
+
+    UC.execute(cmd)
+
+    assert prompts == [
+        "Clean up cached files and unrecognized content? "
+        "This will preserve recognized runtimes."
+    ]
+    assert cached_file.read_bytes() == b"cached"
+    assert unknown_file.read_bytes() == b"unknown"
+
+
+@pytest.mark.parametrize("option", ["purge", "cleanup"])
+def test_global_uninstall_options_reject_tags(fake_config, option):
+    cmd = fake_config
+    cmd.args = ["3.14"]
+    cmd.purge = option == "purge"
+    cmd.cleanup = option == "cleanup"
+
+    with pytest.raises(ArgumentError, match="does not accept runtime tags"):
+        UC.execute(cmd)
+
+
+def test_purge_and_cleanup_are_mutually_exclusive(fake_config):
+    cmd = fake_config
+    cmd.args = []
+    cmd.purge = True
+    cmd.cleanup = True
+
+    with pytest.raises(ArgumentError, match="cannot be combined"):
+        UC.execute(cmd)