Minor refactor, improved test coverage

Author: zooba

src/manage/aliasutils.py

@@ -1,6 +1,7 @@
 import os
 
-from .fsutils import ensure_tree, unlink
+from .exceptions import FilesInUseError
+from .fsutils import atomic_unlink, ensure_tree, unlink
 from .logging import LOGGER
 from .pathutils import Path
 from .tagutils import install_matches_any
@@ -165,24 +166,6 @@ def create_alias(cmd, install, alias, target, *, script_code=None, _link=os.link
             LOGGER.info("Failed to remove %s.", p_script, exc_info=True)
 
 
-def cleanup_alias(cmd):
-    if not cmd.global_dir or not cmd.global_dir.is_dir():
-        return
-
-    alias_written = cmd.scratch.get("aliasutils.create_alias.alias_written") or ()
-
-    for alias in cmd.global_dir.glob("*.exe"):
-        target = alias.with_name(alias.name + ".__target__")
-        script = alias.with_name(alias.name + ".__script__.py")
-        if alias.stem.casefold() not in alias_written:
-            LOGGER.debug("Unlink %s", alias)
-            unlink(alias, f"Attempting to remove {alias} is taking some time. " +
-                           "Ensure it is not is use, and please continue to wait " +
-                           "or press Ctrl+C to abort.")
-            unlink(target)
-            unlink(script)
-
-
 def _parse_entrypoint_line(line):
     line = line.partition("#")[0]
     name, sep, rest = line.partition("=")
@@ -283,32 +266,27 @@ def scan_and_create_entrypoints(cmd, install, shortcut, _create_alias=create_ali
         _create_alias(cmd, install, alias, target, script_code=code)
 
 
-def cleanup_entrypoints(cmd, install_shortcut_pairs):
-    seen_names = set()
-    for install, shortcut in install_shortcut_pairs:
-        for alias, code in _scan(install["prefix"], shortcut.get("dirs")):
-            seen_names.add(alias["name"].casefold())
+def cleanup_alias(cmd, site_dirs_written, *, _unlink_many=atomic_unlink, _scan=_scan):
+    if not cmd.global_dir or not cmd.global_dir.is_dir():
+        return
 
-    # Scan existing aliases
-    scripts = cmd.global_dir.glob("*-script.py")
+    expected = set()
+    for i in cmd.get_installs():
+        expected.update(a.get("name", "").casefold() for a in i.get("alias", ()))
 
-    # Excluding any in seen_names, delete unused aliases
-    for script in scripts:
-        name = script.name.rpartition("-")[0]
-        if name.casefold() in seen_names:
-            continue
+    for i, s in site_dirs_written or ():
+        for alias, code in _scan(i["prefix"], s.get("dirs")):
+            expected.add(alias.get("name", "").casefold())
 
-        alias = cmd.global_dir / (name + ".exe")
-        if not alias.is_file():
+    for alias in cmd.global_dir.glob("*.exe"):
+        if alias.stem.casefold() in expected or alias.name.casefold() in expected:
             continue
-
-        try:
-            unlink(alias)
-            LOGGER.debug("Deleted %s", alias)
-        except OSError:
-            LOGGER.warn("Failed to delete %s", alias)
+        target = alias.with_name(alias.name + ".__target__")
+        script = alias.with_name(alias.name + ".__script__.py")
+        LOGGER.debug("Unlink %s", alias)
         try:
-            unlink(script)
-            LOGGER.debug("Deleted %s", script)
-        except OSError:
-            LOGGER.warn("Failed to delete %s", script)
+            _unlink_many([alias, target, script])
+        except (OSError, FilesInUseError):
+            LOGGER.warn("Failed to remove %s. Ensure it is not in use and run "
+                        "py install --refresh to try again.", alias.name)
+            LOGGER.debug("TRACEBACK", exc_info=True)

src/manage/install_command.py

@@ -254,8 +254,8 @@ def _create_entrypoints(cmd, install, shortcut):
 
 
 def _cleanup_entrypoints(cmd, install_shortcut_pairs):
-    from .aliasutils import cleanup_entrypoints
-    cleanup_entrypoints(cmd, install_shortcut_pairs)
+    # Entry point aliases are cleaned up with regular aliases
+    pass
 
 
 SHORTCUT_HANDLERS = {
@@ -271,6 +271,7 @@ def update_all_shortcuts(cmd, *, _create_alias=None, _cleanup_alias=None):
         from .aliasutils import create_alias as _create_alias
     if not _cleanup_alias:
         from .aliasutils import cleanup_alias as _cleanup_alias
+    from .aliasutils import get_site_dirs
 
     LOGGER.debug("Updating global shortcuts")
     shortcut_written = {}
@@ -309,7 +310,7 @@ def update_all_shortcuts(cmd, *, _create_alias=None, _cleanup_alias=None):
                 create(cmd, i, s)
                 shortcut_written.setdefault(s["kind"], []).append((i, s))
 
-        # Earlier releases may not have site_dirs. If not, assume
+        # Earlier releases may not have site_dirs. If not, assume defaults
         if ("site-dirs" in (cmd.enable_shortcut_kinds or ("site-dirs",)) and
             "site-dirs" not in (cmd.disable_shortcut_kinds or ()) and
             all(s["kind"] != "site-dirs" for s in i.get("shortcuts", ()))):
@@ -319,11 +320,11 @@ def update_all_shortcuts(cmd, *, _create_alias=None, _cleanup_alias=None):
             create(cmd, i, s)
             shortcut_written.setdefault("site-dirs", []).append((i, s))
 
-    _cleanup_alias(cmd)
-
     for k, (_, cleanup) in SHORTCUT_HANDLERS.items():
         cleanup(cmd, shortcut_written.get(k, []))
 
+    _cleanup_alias(cmd, shortcut_written.get("site-dirs", []))
+
 
 def print_cli_shortcuts(cmd):
     if cmd.global_dir and cmd.global_dir.is_dir() and any(cmd.global_dir.glob("*.exe")):

tests/test_alias.py

@@ -354,22 +354,41 @@ def test_scan_entrypoints(tmp_path):
 
 
 def test_cleanup_aliases(fake_config):
+    fake_config.installs = [
+        dict(id="A", alias=[dict(name="A", target="a.exe")], prefix=fake_config.global_dir),
+    ]
+
+    def fake_scan(*a):
+        yield dict(name="B"), "CODE"
+
+    # install/shortcut pairs are irrelevant, since we fake the scan entirely.
+    # It just can't be empty or the scan is skipped.
+    pairs = [
+        (fake_config.installs[0], dict(kind="site-dirs", dirs=[])),
+    ]
+
     root = fake_config.global_dir
     root.mkdir(parents=True, exist_ok=True)
-    (root / "alias1.exe").write_bytes(b"")
-    (root / "alias1.exe.__target__").write_bytes(b"")
-    (root / "alias1.exe.__script__.py").write_bytes(b"")
-    (root / "alias2.exe").write_bytes(b"")
-    (root / "alias2.exe.__target__").write_bytes(b"")
-    (root / "alias2.exe.__script__.py").write_bytes(b"")
-    (root / "alias3.exe").write_bytes(b"")
-    (root / "alias3.exe.__target__").write_bytes(b"")
-    fake_config.scratch["aliasutils.create_alias.alias_written"] = set([
-        "alias1".casefold(),
-        "alias3".casefold(),
-    ])
-    AU.cleanup_alias(fake_config)
-    assert set(f.name for f in root.glob("*")) == set([
-        "alias1.exe", "alias1.exe.__target__", "alias1.exe.__script__.py",
-        "alias3.exe", "alias3.exe.__target__",
-    ])
+    files = ["A.exe", "A.exe.__target__",
+             "B.exe", "B.exe.__script__.py", "B.exe.__target__",
+             "C.exe", "C.exe.__script__.py", "C.exe.__target__"]
+    for f in files:
+        (root / f).write_bytes(b"")
+
+    # Ensure the expect files get requested to be unlinked
+    class Unlinker(list):
+        def __call__(self, names):
+            self.extend(names)
+
+    unlinked = Unlinker()
+    AU.cleanup_alias(fake_config, pairs, _unlink_many=unlinked, _scan=fake_scan)
+    assert set(f.name for f in unlinked) == set(["C.exe", "C.exe.__script__.py", "C.exe.__target__"])
+
+    # Ensure we don't break if unlinking fails
+    def unlink2(names):
+        raise PermissionError("Simulated error")
+    AU.cleanup_alias(fake_config, pairs, _unlink_many=unlink2, _scan=fake_scan)
+
+    # Ensure the actual unlink works
+    AU.cleanup_alias(fake_config, pairs, _scan=fake_scan)
+    assert set(f.name for f in root.glob("*")) == set(files[:-3])