Update the Python module (notably find_ty_bin) for parity with uv (#2852)

Closes https://github.com/astral-sh/ruff/issues/18153

Copy of https://github.com/astral-sh/ruff/pull/23406

python/ty/__init__.py

@@ -0,0 +1,5 @@
+from __future__ import annotations
+
+from ._find_ty import find_ty_bin
+
+__all__ = ["find_ty_bin"]

python/ty/__main__.py

@@ -2,85 +2,26 @@ from __future__ import annotations
 
 import os
 import sys
-import sysconfig
+
+from ty import find_ty_bin
 
 
-def find_ty_bin() -> str:
-    """Return the ty binary path."""
+def _run() -> None:
+    ty = find_ty_bin()
 
-    ty_exe = "ty" + sysconfig.get_config_var("EXE")
-
-    scripts_path = os.path.join(sysconfig.get_path("scripts"), ty_exe)
-    if os.path.isfile(scripts_path):
-        return scripts_path
-
-    if sys.version_info >= (3, 10):
-        user_scheme = sysconfig.get_preferred_scheme("user")
-    elif os.name == "nt":
-        user_scheme = "nt_user"
-    elif sys.platform == "darwin" and sys._framework:
-        user_scheme = "osx_framework_user"
-    else:
-        user_scheme = "posix_user"
-
-    user_path = os.path.join(sysconfig.get_path("scripts", scheme=user_scheme), ty_exe)
-    if os.path.isfile(user_path):
-        return user_path
-
-    # Search in `bin` adjacent to package root (as created by `pip install --target`).
-    pkg_root = os.path.dirname(os.path.dirname(__file__))
-    target_path = os.path.join(pkg_root, "bin", ty_exe)
-    if os.path.isfile(target_path):
-        return target_path
-
-    # Search for pip-specific build environments.
-    #
-    # Expect to find ty in <prefix>/pip-build-env-<rand>/overlay/bin/ty
-    # Expect to find a "normal" folder at <prefix>/pip-build-env-<rand>/normal
-    #
-    # See: https://github.com/pypa/pip/blob/102d8187a1f5a4cd5de7a549fd8a9af34e89a54f/src/pip/_internal/build_env.py#L87
-    paths = os.environ.get("PATH", "").split(os.pathsep)
-    if len(paths) >= 2:
-
-        def get_last_three_path_parts(path: str) -> list[str]:
-            """Return a list of up to the last three parts of a path."""
-            parts = []
-
-            while len(parts) < 3:
-                head, tail = os.path.split(path)
-                if tail or head != path:
-                    parts.append(tail)
-                    path = head
-                else:
-                    parts.append(path)
-                    break
-
-            return parts
-
-        maybe_overlay = get_last_three_path_parts(paths[0])
-        maybe_normal = get_last_three_path_parts(paths[1])
-        if (
-            len(maybe_normal) >= 3
-            and maybe_normal[-1].startswith("pip-build-env-")
-            and maybe_normal[-2] == "normal"
-            and len(maybe_overlay) >= 3
-            and maybe_overlay[-1].startswith("pip-build-env-")
-            and maybe_overlay[-2] == "overlay"
-        ):
-            # The overlay must contain the ty binary.
-            candidate = os.path.join(paths[0], ty_exe)
-            if os.path.isfile(candidate):
-                return candidate
-
-    raise FileNotFoundError(scripts_path)
-
-
-if __name__ == "__main__":
-    ty = os.fsdecode(find_ty_bin())
     if sys.platform == "win32":
         import subprocess
 
-        completed_process = subprocess.run([ty, *sys.argv[1:]])
+        # Avoid emitting a traceback on interrupt
+        try:
+            completed_process = subprocess.run([ty, *sys.argv[1:]])
+        except KeyboardInterrupt:
+            sys.exit(2)
+
         sys.exit(completed_process.returncode)
     else:
         os.execvp(ty, [ty, *sys.argv[1:]])
+
+
+if __name__ == "__main__":
+    _run()

python/ty/_find_ty.py

@@ -0,0 +1,102 @@
+from __future__ import annotations
+
+import os
+import sys
+import sysconfig
+
+
+class TyNotFound(FileNotFoundError): ...
+
+
+def find_ty_bin() -> str:
+    """Return the ty binary path."""
+
+    ty_exe = "ty" + sysconfig.get_config_var("EXE")
+
+    targets = [
+        # The scripts directory for the current Python
+        sysconfig.get_path("scripts"),
+        # The scripts directory for the base prefix
+        sysconfig.get_path("scripts", vars={"base": sys.base_prefix}),
+        # Above the package root, e.g., from `pip install --prefix` or `uv run --with`
+        (
+            # On Windows, with module path `<prefix>/Lib/site-packages/ty`
+            _join(_matching_parents(_module_path(), "Lib/site-packages/ty"), "Scripts")
+            if sys.platform == "win32"
+            # On Unix,  with module path `<prefix>/lib/python3.13/site-packages/ty`
+            else _join(
+                _matching_parents(_module_path(), "lib/python*/site-packages/ty"),
+                "bin",
+            )
+        ),
+        # Adjacent to the package root, e.g., from `pip install --target`
+        # with module path `<target>/ty`
+        _join(_matching_parents(_module_path(), "ty"), "bin"),
+        # The user scheme scripts directory, e.g., `~/.local/bin`
+        sysconfig.get_path("scripts", scheme=_user_scheme()),
+    ]
+
+    seen = []
+    for target in targets:
+        if not target:
+            continue
+        if target in seen:
+            continue
+        seen.append(target)
+        path = os.path.join(target, ty_exe)
+        if os.path.isfile(path):
+            return path
+
+    locations = "\n".join(f" - {target}" for target in seen)
+    raise TyNotFound(
+        f"Could not find the ty binary in any of the following locations:\n{locations}\n"
+    )
+
+
+def _module_path() -> str | None:
+    path = os.path.dirname(__file__)
+    return path
+
+
+def _matching_parents(path: str | None, match: str) -> str | None:
+    """
+    Return the parent directory of `path` after trimming a `match` from the end.
+    The match is expected to contain `/` as a path separator, while the `path`
+    is expected to use the platform's path separator (e.g., `os.sep`). The path
+    components are compared case-insensitively and a `*` wildcard can be used
+    in the `match`.
+    """
+    from fnmatch import fnmatch
+
+    if not path:
+        return None
+    parts = path.split(os.sep)
+    match_parts = match.split("/")
+    if len(parts) < len(match_parts):
+        return None
+
+    if not all(
+        fnmatch(part, match_part)
+        for part, match_part in zip(reversed(parts), reversed(match_parts))
+    ):
+        return None
+
+    return os.sep.join(parts[: -len(match_parts)])
+
+
+def _join(path: str | None, *parts: str) -> str | None:
+    if not path:
+        return None
+    return os.path.join(path, *parts)
+
+
+def _user_scheme() -> str:
+    if sys.version_info >= (3, 10):
+        user_scheme = sysconfig.get_preferred_scheme("user")
+    elif os.name == "nt":
+        user_scheme = "nt_user"
+    elif sys.platform == "darwin" and sys._framework:  # ty: ignore[unresolved-attribute]
+        user_scheme = "osx_framework_user"
+    else:
+        user_scheme = "posix_user"
+    return user_scheme