Create shortcut executables from registered entrypoints

.github/workflows/build.yml

@@ -172,6 +172,25 @@ jobs:
         PYTHON_MANAGER_CONFIG: .\test-config.json
         PYMANAGER_DEBUG: true
 
+    - name: 'Validate entrypoint script'
+      run: |
+        $env:PYTHON_MANAGER_CONFIG = (gi $env:PYTHON_MANAGER_CONFIG).FullName
+        cd .\test_installs\_bin
+        del pip* -Verbose
+        pymanager install --refresh
+        dir pip*
+        Get-Item .\pip.exe
+        Get-Item .\pip.exe.__target__
+        Get-Content .\pip.exe.__target__
+        Get-Item .\pip.exe.__script__.py
+        Get-Content .\pip.exe.__script__.py
+        .\pip.exe --version
+      env:
+        PYTHON_MANAGER_INCLUDE_UNMANAGED: false
+        PYTHON_MANAGER_CONFIG: .\test-config.json
+        PYMANAGER_DEBUG: true
+      shell: powershell
+
     - name: 'Offline bundle download and install'
       run: |
         pymanager list --online 3 3-32 3-64 3-arm64

src/manage/aliasutils.py

@@ -0,0 +1,371 @@
+import os
+
+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.
+# Be careful to use the most compatible Python code possible.
+try:
+    if sys.path[0]:
+        if sys.argv[0].startswith(sys.path[0]):
+            sys.path[0] = ""
+        else:
+            open(sys.path[0] + "/" + sys.argv[0], "rb").close()
+            sys.path[0] = ""
+except OSError:
+    pass
+except AttributeError:
+    pass
+except IndexError:
+    pass
+
+# Replace argv[0] with our executable instead of the script name.
+try:
+    if sys.argv[0][-14:].upper() == ".__SCRIPT__.PY":
+        sys.argv[0] = sys.argv[0][:-14]
+        sys.orig_argv[0] = sys.argv[0]
+except AttributeError:
+    pass
+except IndexError:
+    pass
+
+from {mod} import {func}
+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:
+            if not all(s.isidentifier() for s in self.mod.split(".")):
+                LOGGER.warn("Alias %s has an entrypoint with invalid module "
+                            "%r.", self.name, self.mod)
+                return None
+            if not all(s.isidentifier() for s in self.func.split(".")):
+                LOGGER.warn("Alias %s has an entrypoint with invalid function "
+                            "%r.", self.name, self.func)
+                return None
+            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}"
+    if plat_launcher.is_file():
+        return plat_launcher
+    return launcher
+
+
+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 windowed:
+        launcher = cmd.launcherw_exe or launcher
+
+    if plat:
+        LOGGER.debug("Checking for launcher for platform -%s", plat)
+        launcher = _if_exists(launcher, f"-{plat}")
+    if not launcher.is_file():
+        LOGGER.debug("Checking for launcher for default platform %s", cmd.default_platform)
+        launcher = _if_exists(launcher, cmd.default_platform)
+    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", name, target, launcher)
+    if not launcher or not launcher.is_file():
+        raise NoLauncherTemplateError()
+
+    try:
+        launcher_bytes = launcher.read_bytes()
+    except OSError:
+        warnings_shown = cmd.scratch.setdefault("aliasutils.create_alias.warnings_shown", set())
+        if str(launcher) not in warnings_shown:
+            LOGGER.warn("Failed to read launcher template at %s.", launcher)
+            warnings_shown.add(str(launcher))
+        LOGGER.debug("Failed to read %s", launcher, exc_info=True)
+        return
+
+    existing_bytes = b''
+    try:
+        with open(p, 'rb') as f:
+            existing_bytes = f.read(len(launcher_bytes) + 1)
+    except FileNotFoundError:
+        pass
+    except OSError:
+        LOGGER.debug("Failed to read existing alias launcher.")
+
+    launcher_remap = cmd.scratch.setdefault("aliasutils.create_alias.launcher_remap", {})
+    if existing_bytes == launcher_bytes:
+        # Valid existing launcher, so save its path in case we need it later
+        # for a hard link.
+        launcher_remap.setdefault(launcher.name, p)
+    else:
+        # First try and create a hard link
+        unlink(p)
+        try:
+            _link(launcher, p)
+            LOGGER.debug("Created %s as hard link to %s", p.name, launcher.name)
+        except OSError as ex:
+            if ex.winerror != 17:
+                # Report errors other than cross-drive links
+                LOGGER.debug("Failed to create hard link for command.", exc_info=True)
+            launcher2 = launcher_remap.get(launcher.name)
+            if launcher2:
+                try:
+                    _link(launcher2, p)
+                    LOGGER.debug("Created %s as hard link to %s", p.name, launcher2.name)
+                except FileNotFoundError:
+                    raise
+                except OSError:
+                    LOGGER.debug("Failed to create hard link to fallback launcher")
+                    launcher2 = None
+            if not launcher2:
+                try:
+                    p.write_bytes(launcher_bytes)
+                    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.", name)
+                    LOGGER.debug("TRACEBACK", exc_info=True)
+
+    p_target = p.with_name(p.name + ".__target__")
+    do_update = True
+    try:
+        do_update = not target.match(p_target.read_text(encoding="utf-8"))
+    except FileNotFoundError:
+        pass
+    except (OSError, UnicodeDecodeError):
+        LOGGER.debug("Failed to read existing target path.", exc_info=True)
+
+    if do_update:
+        p_target.write_text(str(target), encoding="utf-8")
+
+    p_script = p.with_name(p.name + ".__script__.py")
+    if script_code:
+        do_update = True
+        try:
+            do_update = p_script.read_text(encoding="utf-8") != script_code
+        except FileNotFoundError:
+            pass
+        except (OSError, UnicodeDecodeError):
+            LOGGER.debug("Failed to read existing script file.", exc_info=True)
+        if do_update:
+            p_script.write_text(script_code, encoding="utf-8")
+    else:
+        try:
+            unlink(p_script)
+        except OSError:
+            LOGGER.error("Failed to clean up existing alias. Re-run with -v "
+                         "or check the install log for details.")
+            LOGGER.info("Failed to remove %s.", p_script)
+            LOGGER.debug("TRACEBACK", exc_info=True)
+
+
+def _parse_entrypoint_line(line):
+    line = line.partition("#")[0]
+    name, sep, rest = line.partition("=")
+    name = name.strip()
+    if name and name[0].isalnum() and sep and rest:
+        mod, sep, rest = rest.partition(":")
+        mod = mod.strip()
+        if mod and sep and rest:
+            func, sep, extra = rest.partition("[")
+            func = func.strip()
+            if func:
+                return name, mod, func
+    return None, None, None
+
+
+def _readlines(path):
+    try:
+        f = open(path, "r", encoding="utf-8", errors="strict")
+    except OSError:
+        LOGGER.debug("Failed to read %s", path, exc_info=True)
+        return
+
+    with f:
+        try:
+            while True:
+                yield next(f)
+        except StopIteration:
+            return
+        except UnicodeDecodeError:
+            LOGGER.debug("Failed to decode contents of %s", path, exc_info=True)
+            return
+
+
+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()]
+    entrypoints = [f for f in [d / "entry_points.txt" for d in dist_info] if f.is_file()]
+    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:
+        alias = None
+        for line in _readlines(ep):
+            if line.strip() == "[console_scripts]":
+                alias = dict(windowed=0)
+            elif line.strip() == "[gui_scripts]":
+                alias = dict(windowed=1)
+            elif line.lstrip().startswith("["):
+                alias = None
+            elif alias is not None:
+                name, mod, func = _parse_entrypoint_line(line)
+                if name and mod and func:
+                    yield AliasInfo(install=install, name=name,
+                                    mod=mod, func=func, **alias)
+
+
+def _scan(install, prefix, dirs):
+    for dirname in dirs or ():
+        root = prefix / dirname
+        yield from _scan_one(install, root)
+
+
+def calculate_aliases(cmd, install, *, _scan=_scan):
+    LOGGER.debug("Calculating aliases for %s", install["id"])
+
+    prefix = install["prefix"]
+
+    default_alias = None
+    default_alias_w = None
+
+    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
+
+    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
+        written.add(n)
+
+        if not alias.target:
+            LOGGER.debug("No suitable alias found for %s. Skipping", alias.name)
+            continue
+
+        target = alias.install["prefix"] / alias.target
+        try:
+            _create_alias(
+                cmd,
+                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 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.name.casefold() in expected:
+            continue
+        target = alias.with_name(alias.name + ".__target__")
+        script = alias.with_name(alias.name + ".__script__.py")
+        LOGGER.debug("Unlink %s", alias)
+        try:
+            _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)