Improved script and fixed names

Author: zooba

src/manage/aliasutils.py

@@ -5,6 +5,37 @@
 from .pathutils import Path
 from .tagutils import install_matches_any
 
+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}())
+"""
 
 def _if_exists(launcher, plat):
     suffix = "." + launcher.suffix.lstrip(".")
@@ -15,13 +46,22 @@ def _if_exists(launcher, plat):
 
 
 def create_alias(cmd, install, alias, target, *, script_code=None, _link=os.link):
-    p = (cmd.global_dir / alias["name"])
+    p = cmd.global_dir / alias["name"]
+    if not p.match("*.exe"):
+        p = p.with_name(p.name + ".exe")
     target = Path(target)
     ensure_tree(p)
     launcher = cmd.launcher_exe
     if alias.get("windowed"):
         launcher = cmd.launcherw_exe or launcher
 
+    alias_written = cmd.scratch.setdefault("aliasutils.create_alias.alias_written", set())
+    n = p.stem.casefold()
+    if n in alias_written:
+        # We've already written this alias in this session, so skip it.
+        return
+    alias_written.add(n)
+
     plat = install["tag"].rpartition("-")[-1]
     if plat:
         LOGGER.debug("Checking for launcher for platform -%s", plat)
@@ -60,7 +100,6 @@ def create_alias(cmd, install, alias, target, *, script_code=None, _link=os.link
         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.
@@ -106,7 +145,7 @@ def create_alias(cmd, install, alias, target, *, script_code=None, _link=os.link
     if do_update:
         p_target.write_text(str(target), encoding="utf-8")
 
-    p_script = p.with_name(p.name + "-script.py")
+    p_script = p.with_name(p.name + ".__script__.py")
     if script_code:
         do_update = True
         try:
@@ -126,6 +165,24 @@ 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("=")
@@ -170,7 +227,7 @@ def _scan_one(root):
                     if name and mod and func:
                         yield (
                             {**alias, "name": name},
-                            f"import sys; from {mod} import {func}; sys.exit({func}())",
+                            SCRIPT_CODE.format(mod=mod, func=func),
                         )
 
 

src/manage/install_command.py

@@ -266,12 +266,12 @@ def _cleanup_entrypoints(cmd, install_shortcut_pairs):
 }
 
 
-def update_all_shortcuts(cmd, *, _create_alias=None):
+def update_all_shortcuts(cmd, *, _create_alias=None, _cleanup_alias=None):
     if not _create_alias:
         from .aliasutils import create_alias as _create_alias
+        from .aliasutils import cleanup_alias as _cleanup_alias
 
     LOGGER.debug("Updating global shortcuts")
-    alias_written = set()
     shortcut_written = {}
     for i in cmd.get_installs():
         if cmd.global_dir:
@@ -288,14 +288,11 @@ def update_all_shortcuts(cmd, *, _create_alias=None):
                     aliases.append({**alias_2[0], "name": "pythonw.exe"})
 
             for a in aliases:
-                if a["name"].casefold() in alias_written:
-                    continue
                 target = i["prefix"] / a["target"]
                 if not target.is_file():
                     LOGGER.warn("Skipping alias '%s' because target '%s' does not exist", a["name"], a["target"])
                     continue
                 _create_alias(cmd, i, a, target)
-                alias_written.add(a["name"].casefold())
 
         for s in i.get("shortcuts", ()):
             if cmd.enable_shortcut_kinds and s["kind"] not in cmd.enable_shortcut_kinds:
@@ -321,15 +318,7 @@ def update_all_shortcuts(cmd, *, _create_alias=None):
             create(cmd, i, s)
             shortcut_written.setdefault("site-dirs", []).append((i, s))
 
-    if cmd.global_dir and cmd.global_dir.is_dir() and cmd.launcher_exe:
-        for target in cmd.global_dir.glob("*.exe.__target__"):
-            alias = target.with_suffix("")
-            if alias.name.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.")
-                target.unlink()
+    _cleanup_alias(cmd)
 
     for k, (_, cleanup) in SHORTCUT_HANDLERS.items():
         cleanup(cmd, shortcut_written.get(k, []))

src/manage/pathutils.py

@@ -44,11 +44,17 @@ def __bool__(self):
 
     @property
     def stem(self):
-        return self.name.rpartition(".")[0]
+        stem, dot, suffix = self.name.rpartition(".")
+        if not dot:
+            return suffix
+        return stem
 
     @property
     def suffix(self):
-        return self.name.rpartition(".")[2]
+        stem, dot, suffix = self.name.rpartition(".")
+        if not dot:
+            return ""
+        return dot + suffix
 
     @property
     def parent(self):

src/pymanager/launcher.cpp

@@ -145,43 +145,56 @@ get_executable(wchar_t *executable, unsigned int bufferSize)
 int
 insert_script(int *argc, wchar_t ***argv)
 {
-    DWORD len = GetModuleFileNameW(NULL, NULL, 0);
-    if (len == 0) {
-        return HRESULT_FROM_WIN32(GetLastError());
-    } else if (len < 5) {
-        return 0;
-    }
-
     HANDLE ph = GetProcessHeap();
-    DWORD path_len = len + 7;
-    wchar_t *path = (wchar_t *)HeapAlloc(ph, HEAP_ZERO_MEMORY, sizeof(wchar_t) * path_len);
-    len = path ? GetModuleFileNameW(NULL, path, path_len) : 0;
-    if (len == 0) {
-        return HRESULT_FROM_WIN32(GetLastError());
-    }
+    wchar_t *path = NULL;
+    DWORD path_len = 0;
+    DWORD len = 0;
+    int error = 0;
+    const wchar_t *SUFFIX = L".__script__.py";
+
+    // Get our path in a dynamic buffer with enough space to add SUFFIX
+    while (len >= path_len) {
+        if (path) {
+            HeapFree(ph, 0, path);
+        }
+        path_len += 260;
 
-    if (wcsicmp(&path[len - 4], L".exe")) {
-        HeapFree(ph, 0, path);
-        return 0;
+        path = (wchar_t *)HeapAlloc(ph, HEAP_ZERO_MEMORY, sizeof(wchar_t) * path_len);
+        if (!path) {
+            return HRESULT_FROM_WIN32(GetLastError());
+        }
+
+        len = GetModuleFileNameW(NULL, path, path_len - wcslen(SUFFIX));
+        if (len == 0) {
+            error = GetLastError();
+            HeapFree(ph, 0, path);
+            return HRESULT_FROM_WIN32(error);
+        }
     }
 
-    wcscpy_s(&path[len - 4], path_len, L"-script.py");
+    wcscpy_s(&path[len], path_len, SUFFIX);
+
+    // Check that we have a script file. FindFirstFile should be fastest.
     WIN32_FIND_DATAW fd;
     HANDLE fh = FindFirstFileW(path, &fd);
     if (fh == INVALID_HANDLE_VALUE) {
-        int err = GetLastError();
+        error = GetLastError();
         HeapFree(ph, 0, path);
-        switch (err) {
+        switch (error) {
         case ERROR_INVALID_FUNCTION:
         case ERROR_FILE_NOT_FOUND:
         case ERROR_PATH_NOT_FOUND:
+            // This is the typical exit for normal launches. We ought to be nice
+            // and fast up until this point, but can be slower through every
+            // other path.
             return 0;
         default:
-            return HRESULT_FROM_WIN32(GetLastError());
+            return HRESULT_FROM_WIN32(error);
         }
     }
     CloseHandle(fh);
 
+    // Create a new argv that will be used to launch the script.
     wchar_t **argv2 = (wchar_t **)HeapAlloc(ph, HEAP_ZERO_MEMORY, sizeof(wchar_t *) * (*argc + 1));
     if (!argv2) {
         HeapFree(ph, 0, path);

tests/test_pathutils.py

@@ -15,3 +15,18 @@ def test_path_match():
     assert not p.match("example*")
     assert not p.match("example*.com")
     assert not p.match("*ple*")
+
+
+def test_path_stem():
+    p = Path("python3.12.exe")
+    assert p.stem == "python3.12"
+    assert p.suffix == ".exe"
+    p = Path("python3.12")
+    assert p.stem == "python3"
+    assert p.suffix == ".12"
+    p = Path("python3")
+    assert p.stem == "python3"
+    assert p.suffix == ""
+    p = Path(".exe")
+    assert p.stem == ""
+    assert p.suffix == ".exe"