diff --git a/python/ty/__main__.py b/python/ty/__main__.py index 47c44e37a..255b7f5cb 100644 --- a/python/ty/__main__.py +++ b/python/ty/__main__.py @@ -75,12 +75,49 @@ def get_last_three_path_parts(path: str) -> list[str]: raise FileNotFoundError(scripts_path) -if __name__ == "__main__": +def _detect_virtualenv() -> str: + """ + Find the virtual environment path for the current Python executable. + """ + + # If it's already set, then just use it + value = os.getenv("VIRTUAL_ENV") + if value: + return value + + # Otherwise, check if we're in a venv + venv_marker = os.path.join(sys.prefix, "pyvenv.cfg") + + if os.path.exists(venv_marker): + return sys.prefix + + return "" + + +def _run() -> None: ty = os.fsdecode(find_ty_bin()) + + env = os.environ.copy() + venv = _detect_virtualenv() + if venv: + env.setdefault("VIRTUAL_ENV", venv) + + # Let `ty` know that it was spawned by this Python interpreter + env["TY__PARENT_INTERPRETER"] = sys.executable + 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:]], env=env) + except KeyboardInterrupt: + sys.exit(2) + sys.exit(completed_process.returncode) else: - os.execvp(ty, [ty, *sys.argv[1:]]) + os.execvpe(ty, [ty, *sys.argv[1:]], env=env) + + +if __name__ == "__main__": + _run()