From 12d5f1fdc6bf541b1fcf4f5570bd814c139b5b9c Mon Sep 17 00:00:00 2001
From: Steve Dower <steve.dower@python.org>
Date: Thu, 15 May 2025 00:00:39 +0100
Subject: [PATCH] Improves purging and shortcut removal. Fixes #102

---
 src/manage/pep514utils.py       |  7 +++++--
 src/manage/uninstall_command.py | 19 +++++++++----------
 2 files changed, 14 insertions(+), 12 deletions(-)

diff --git a/src/manage/pep514utils.py b/src/manage/pep514utils.py
index ab011b4..91ccb72 100644
--- a/src/manage/pep514utils.py
+++ b/src/manage/pep514utils.py
@@ -256,20 +256,23 @@ def update_registry(root_name, install, data, warn_for=[]):
 
 
 def cleanup_registry(root_name, keep, warn_for=[]):
+    LOGGER.debug("Cleaning up registry entries")
     hive, name = _split_root(root_name)
     with _reg_open(hive, name, writable=True) as root:
-        for company_name in _iter_keys(root):
+        for company_name in list(_iter_keys(root)):
             any_left = False
             with winreg.OpenKey(root, company_name, access=winreg.KEY_ALL_ACCESS) as company:
-                for tag_name in _iter_keys(company):
+                for tag_name in list(_iter_keys(company)):
                     # Calculate whether to show warnings or not
                     install = {"company": company_name, "tag": tag_name}
                     allow_warn = install_matches_any(install, warn_for)
 
                     if (f"{company_name}\\{tag_name}" in keep
                         or not _is_tag_managed(company, tag_name, allow_warn=allow_warn)):
+                        LOGGER.debug("Skipping %s\\%s\\%s", root_name, company_name, tag_name)
                         any_left = True
                     else:
+                        LOGGER.debug("Removing %s\\%s\\%s", root_name, company_name, tag_name)
                         _reg_rmtree(company, tag_name)
             if not any_left:
                 _delete_key(root, company_name)
diff --git a/src/manage/uninstall_command.py b/src/manage/uninstall_command.py
index 9574ef6..def69f9 100644
--- a/src/manage/uninstall_command.py
+++ b/src/manage/uninstall_command.py
@@ -1,7 +1,7 @@
 from .exceptions import ArgumentError, FilesInUseError
 from .fsutils import rmtree, unlink
 from .installs import get_matching_install_tags
-from .install_command import update_all_shortcuts
+from .install_command import SHORTCUT_HANDLERS, update_all_shortcuts
 from .logging import LOGGER
 from .pathutils import PurePath
 from .tagutils import tag_or_range
@@ -9,6 +9,8 @@
 
 def _iterdir(p, only_files=False):
     try:
+        if only_files:
+            return [f for f in p.iterdir() if p.is_file()]
         return list(p.iterdir())
     except FileNotFoundError:
         LOGGER.debug("Skipping %s because it does not exist", p)
@@ -42,18 +44,15 @@ def execute(cmd):
                     LOGGER.warn("Unable to purge %s because it is still in use.",
                                 i["display-name"])
                     continue
-            LOGGER.info("Purging saved downloads")
-            for f in _iterdir(cmd.install_dir):
-                LOGGER.debug("Purging %s", f)
-                try:
-                    rmtree(f, after_5s_warning=warn_msg.format("cached downloads"),
-                           remove_ext_first=("exe", "dll", "json"))
-                except FilesInUseError:
-                    pass
-            LOGGER.info("Purging global commands")
+            LOGGER.info("Purging saved downloads from %s", cmd.download_dir)
+            rmtree(cmd.download_dir, after_5s_warning=warn_msg.format("cached downloads"))
+            LOGGER.info("Purging global commands from %s", cmd.global_dir)
             for f in _iterdir(cmd.global_dir):
                 LOGGER.debug("Purging %s", f)
                 rmtree(f, after_5s_warning=warn_msg.format("global commands"))
+            LOGGER.info("Purging all shortcuts")
+            for _, cleanup in SHORTCUT_HANDLERS.values():
+                cleanup(cmd, [])
         LOGGER.debug("END uninstall_command.execute")
         return