Refactor and simplify code for aliases

Author: zooba

src/manage/aliasutils.py

@@ -1,11 +1,15 @@
 import os
 
-from .exceptions import FilesInUseError
+from .exceptions import FilesInUseError, NoLauncherTemplateError
 from .fsutils import atomic_unlink, ensure_tree, unlink
 from .logging import LOGGER
 from .pathutils import Path
 from .tagutils import install_matches_any
 
+_EXE = ".exe".casefold()
+
+DEFAULT_SITE_DIRS = ["Lib\\site-packages", "Scripts"]
+
 SCRIPT_CODE = """import sys
 
 # Clear sys.path[0] if it contains this script.
@@ -38,6 +42,33 @@
 sys.exit({func}())
 """
 
+
+class AliasInfo:
+    def __init__(self, **kwargs):
+        self.install = kwargs.get("install")
+        self.name = kwargs.get("name")
+        self.windowed = kwargs.get("windowed", 0)
+        self.target = kwargs.get("target")
+        self.mod = kwargs.get("mod")
+        self.func = kwargs.get("func")
+
+    def replace(self, **kwargs):
+        return AliasInfo(**{
+            "install": self.install,
+            "name": self.name,
+            "windowed": self.windowed,
+            "target": self.target,
+            "mod": self.mod,
+            "func": self.func,
+            **kwargs,
+        })
+
+    @property
+    def script_code(self):
+        if self.mod and self.func:
+            return SCRIPT_CODE.format(mod=self.mod, func=self.func)
+
+
 def _if_exists(launcher, plat):
     suffix = "." + launcher.suffix.lstrip(".")
     plat_launcher = launcher.parent / f"{launcher.stem}{plat}{suffix}"
@@ -46,24 +77,17 @@ def _if_exists(launcher, plat):
     return launcher
 
 
-def create_alias(cmd, install, alias, target, aliases_written, *, script_code=None, _link=os.link):
-    p = cmd.global_dir / alias["name"]
+def _create_alias(cmd, *, name, target, plat=None, windowed=0, script_code=None, _link=os.link):
+    p = cmd.global_dir / name
     if not p.match("*.exe"):
         p = p.with_name(p.name + ".exe")
     if not isinstance(target, Path):
         target = Path(target)
     ensure_tree(p)
     launcher = cmd.launcher_exe
-    if alias.get("windowed"):
+    if windowed:
         launcher = cmd.launcherw_exe or launcher
 
-    n = p.stem.casefold()
-    if n in aliases_written:
-        # We've already written this alias in this session, so skip it.
-        return
-    aliases_written.add(n)
-
-    plat = install["tag"].rpartition("-")[-1]
     if plat:
         LOGGER.debug("Checking for launcher for platform -%s", plat)
         launcher = _if_exists(launcher, f"-{plat}")
@@ -73,13 +97,9 @@ def create_alias(cmd, install, alias, target, aliases_written, *, script_code=No
     if not launcher.is_file():
         LOGGER.debug("Checking for launcher for -64")
         launcher = _if_exists(launcher, "-64")
-    LOGGER.debug("Create %s linking to %s using %s", alias["name"], target, launcher)
+    LOGGER.debug("Create %s linking to %s using %s", name, target, launcher)
     if not launcher or not launcher.is_file():
-        if install_matches_any(install, getattr(cmd, "tags", None)):
-            LOGGER.warn("Skipping %s alias because the launcher template was not found.", alias["name"])
-        else:
-            LOGGER.debug("Skipping %s alias because the launcher template was not found.", alias["name"])
-        return
+        raise NoLauncherTemplateError()
 
     try:
         launcher_bytes = launcher.read_bytes()
@@ -131,7 +151,7 @@ def create_alias(cmd, install, alias, target, aliases_written, *, script_code=No
                     LOGGER.debug("Created %s as copy of %s", p.name, launcher.name)
                     launcher_remap[launcher.name] = p
                 except OSError:
-                    LOGGER.error("Failed to create global command %s.", alias["name"])
+                    LOGGER.error("Failed to create global command %s.", name)
                     LOGGER.debug("TRACEBACK", exc_info=True)
 
     p_target = p.with_name(p.name + ".__target__")
@@ -199,12 +219,13 @@ def _readlines(path):
             return
 
 
-def _scan_one(root):
+def _scan_one(install, root):
     # Scan d for dist-info directories with entry_points.txt
     dist_info = [d for d in root.glob("*.dist-info") if d.is_dir()]
-    LOGGER.debug("Found %i dist-info directories in %s", len(dist_info), root)
     entrypoints = [f for f in [d / "entry_points.txt" for d in dist_info] if f.is_file()]
-    LOGGER.debug("Found %i entry_points.txt files in %s", len(entrypoints), root)
+    if len(entrypoints):
+        LOGGER.debug("Found %i entry_points.txt files in %i dist-info in %s",
+                     len(entrypoints), len(dist_info), root)
 
     # Filter down to [console_scripts] and [gui_scripts]
     for ep in entrypoints:
@@ -219,64 +240,117 @@ def _scan_one(root):
             elif alias is not None:
                 name, mod, func = _parse_entrypoint_line(line)
                 if name and mod and func:
-                    yield (
-                        {**alias, "name": name},
-                        SCRIPT_CODE.format(mod=mod, func=func),
-                    )
+                    yield AliasInfo(install=install, name=name,
+                                    mod=mod, func=func, **alias)
 
 
-def _scan(prefix, dirs):
+def _scan(install, prefix, dirs):
     for dirname in dirs or ():
         root = prefix / dirname
-        yield from _scan_one(root)
+        yield from _scan_one(install, root)
 
 
-def scan_and_create_entrypoints(cmd, install, shortcut, aliases_written, *, _create_alias=create_alias, _scan=_scan):
-    prefix = install["prefix"]
+def calculate_aliases(cmd, install, *, _scan=_scan):
+    LOGGER.debug("Calculating aliases for %s", install["id"])
 
-    aliases = list(install.get("alias", ()))
-    alias_1 = [a for a in aliases if not a.get("windowed")]
-    # If no windowed targets, we'll use the non-windowed one
-    alias_2 = [a for a in aliases if a.get("windowed")] or alias_1
+    prefix = install["prefix"]
 
-    targets = [
-        (prefix / alias_1[0]["target"]) if alias_1 else None,
-        (prefix / alias_2[0]["target"]) if alias_2 else None,
-    ]
+    default_alias = None
+    default_alias_w = None
 
-    if not any(targets):
-        LOGGER.debug("No suitable alias found for %s. Skipping entrypoints",
-                     install["id"])
+    for a in install.get("alias", ()):
+        target = prefix / a["target"]
+        if not target.is_file():
+            LOGGER.warn("Skipping alias '%s' because target '%s' does not exist",
+                        a["name"], a["target"])
+            continue
+        ai = AliasInfo(install=install, **a)
+        yield ai
+        if a.get("windowed") and not default_alias_w:
+            default_alias_w = ai
+        if not default_alias:
+            default_alias = ai
+
+    if not default_alias_w:
+        default_alias_w = default_alias
+
+    if install.get("default"):
+        if default_alias:
+            yield default_alias.replace(name="python")
+        if default_alias_w:
+            yield default_alias_w.replace(name="pythonw", windowed=1)
+
+    site_dirs = DEFAULT_SITE_DIRS
+    for s in install.get("shortcuts", ()):
+        if s.get("kind") == "site-dirs":
+            site_dirs = s.get("dirs", ())
+            break
+
+    for ai in _scan(install, prefix, site_dirs):
+        if ai.windowed and default_alias_w:
+            yield ai.replace(target=default_alias_w.target)
+        elif not ai.windowed and default_alias:
+            yield ai.replace(target=default_alias.target)
+
+
+def create_aliases(cmd, aliases, *, _create_alias=_create_alias):
+    if not cmd.global_dir:
         return
 
-    for alias, code in _scan(prefix, shortcut.get("dirs")):
-        # Copy the launcher template and create a standard __target__ file
-        target = targets[1 if alias.get("windowed", 0) else 0]
-        if not target:
-            LOGGER.debug("No suitable alias found for %s. Skipping this " +
-                         "entrypoint", alias["name"])
+    written = set()
+
+    LOGGER.debug("Creating aliases")
+
+    for alias in aliases:
+        if not alias.name:
+            LOGGER.debug("Invalid alias info provided with no name.")
+            continue
+
+        n = alias.name.casefold().removesuffix(_EXE)
+        if n in written:
+            # We've already written this alias, so skip it.
             continue
-        _create_alias(cmd, install, alias, target, aliases_written, script_code=code)
+        written.add(n)
 
+        if not alias.target:
+            LOGGER.debug("No suitable alias found for %s. Skipping", alias.name)
+            continue
 
-def cleanup_alias(cmd, site_dirs_written, *, _unlink_many=atomic_unlink, _scan=_scan):
+        target = alias.install["prefix"] / alias.target
+        try:
+            _create_alias(
+                cmd,
+                install=alias.install,
+                name=alias.name,
+                plat=alias.install.get("tag", "").rpartition("-")[2],
+                target=target,
+                script_code=alias.script_code,
+                windowed=alias.windowed,
+            )
+        except NoLauncherTemplateError:
+            if install_matches_any(alias.install, getattr(cmd, "tags", None)):
+                LOGGER.warn("Skipping %s alias because "
+                            "the launcher template was not found.", alias.name)
+            else:
+                LOGGER.debug("Skipping %s alias because "
+                             "the launcher template was not found.", alias.name)
+
+
+
+def cleanup_aliases(cmd, *, preserve, _unlink_many=atomic_unlink):
     if not cmd.global_dir or not cmd.global_dir.is_dir():
         return
 
+    LOGGER.debug("Cleaning up aliases")
     expected = set()
-    for i in cmd.get_installs():
-        expected.update(a.get("name", "").casefold() for a in i.get("alias", ()))
-
-    if expected:
-        expected.add("python".casefold())
-        expected.add("pythonw".casefold())
-
-    for i, s in site_dirs_written or ():
-        for alias, code in _scan(i["prefix"], s.get("dirs")):
-            expected.add(alias.get("name", "").casefold())
+    for alias in preserve:
+        if alias.name:
+            n = alias.name.casefold().removesuffix(_EXE) + _EXE
+            expected.add(n)
 
+    LOGGER.debug("Retaining %d aliases", len(expected))
     for alias in cmd.global_dir.glob("*.exe"):
-        if alias.stem.casefold() in expected or alias.name.casefold() in expected:
+        if alias.name.casefold() in expected:
             continue
         target = alias.with_name(alias.name + ".__target__")
         script = alias.with_name(alias.name + ".__script__.py")

src/manage/arputils.py

@@ -25,7 +25,7 @@ def _self_cmd():
     if not appdata:
         appdata = os.path.expanduser(r"~\AppData\Local")
     apps = Path(appdata) / r"Microsoft\WindowsApps"
-    LOGGER.debug("Searching %s for pymanager.exe", apps)
+    LOGGER.debug("Searching %s for pymanager.exe for ARP entries", apps)
     for d in apps.iterdir():
         if not d.match("PythonSoftwareFoundation.PythonManager_*"):
             continue

src/manage/exceptions.py

@@ -75,3 +75,8 @@ def __init__(self):
 class FilesInUseError(Exception):
     def __init__(self, files):
         self.files = files
+
+
+class NoLauncherTemplateError(Exception):
+    def __init__(self):
+        super().__init__("No suitable launcher template was found.")