From 26b3c69ee4308be9a8d20b5dc368ff77f8dcc7de Mon Sep 17 00:00:00 2001
From: Sebastian Rittau <srittau@rittau.biz>
Date: Thu, 3 Oct 2024 04:48:52 +0200
Subject: [PATCH] Add option to keep tmp dir to stubtest script (#12151)

---
 tests/_utils.py               |  4 ++++
 tests/stubtest_third_party.py | 31 +++++++++++++++++++++++++++----
 tests/typecheck_typeshed.py   |  5 +++--
 3 files changed, 34 insertions(+), 6 deletions(-)

diff --git a/tests/_utils.py b/tests/_utils.py
index 2879ee486..7381707ed 100644
--- a/tests/_utils.py
+++ b/tests/_utils.py
@@ -47,6 +47,10 @@ def print_command(cmd: str | Iterable[str]) -> None:
     print(colored(f"Running: {cmd}", "blue"))
 
 
+def print_info(message: str) -> None:
+    print(colored(message, "blue"))
+
+
 def print_error(error: str, end: str = "\n", fix_path: tuple[str, str] = ("", "")) -> None:
     error_split = error.split("\n")
     old, new = fix_path
diff --git a/tests/stubtest_third_party.py b/tests/stubtest_third_party.py
index c22c81732..f2fc317e9 100755
--- a/tests/stubtest_third_party.py
+++ b/tests/stubtest_third_party.py
@@ -9,6 +9,7 @@ import subprocess
 import sys
 import tempfile
 from pathlib import Path
+from shutil import rmtree
 from textwrap import dedent
 from typing import NoReturn
 
@@ -21,13 +22,19 @@ from _utils import (
     get_mypy_req,
     print_divider,
     print_error,
+    print_info,
     print_success_msg,
     tests_path,
 )
 
 
 def run_stubtest(
-    dist: Path, *, parser: argparse.ArgumentParser, verbose: bool = False, specified_platforms_only: bool = False
+    dist: Path,
+    *,
+    parser: argparse.ArgumentParser,
+    verbose: bool = False,
+    specified_platforms_only: bool = False,
+    keep_tmp_dir: bool = False,
 ) -> bool:
     dist_name = dist.name
     try:
@@ -51,8 +58,9 @@ def run_stubtest(
         print(colored(f"skipping (requires Python {metadata.requires_python})", "yellow"))
         return True
 
-    with tempfile.TemporaryDirectory() as tmp:
-        venv_dir = Path(tmp)
+    tmp = tempfile.mkdtemp(prefix="stubtest-")  # TODO: Python 3.12: Use TemporaryDirectory
+    venv_dir = Path(tmp)
+    try:
         try:
             subprocess.run(["uv", "venv", venv_dir, "--seed"], capture_output=True, check=True)
         except subprocess.CalledProcessError as e:
@@ -150,9 +158,12 @@ def run_stubtest(
             print("Python version: ", end="", flush=True)
             ret = subprocess.run([sys.executable, "-VV"], capture_output=True)
             print_command_output(ret)
+
             print("\nRan with the following environment:")
             ret = subprocess.run([pip_exe, "freeze", "--all"], capture_output=True)
             print_command_output(ret)
+            if keep_tmp_dir:
+                print("Path to virtual environment:", venv_dir, flush=True)
 
             print_divider()
             main_allowlist_path = allowlists_path(dist_name) / "stubtest_allowlist.txt"
@@ -173,6 +184,11 @@ def run_stubtest(
             return False
         else:
             print_success_msg()
+            if keep_tmp_dir:
+                print_info(f"Virtual environment kept at: {venv_dir}")
+    finally:
+        if not keep_tmp_dir:
+            rmtree(venv_dir)
 
     if verbose:
         print_commands(dist, pip_cmd, stubtest_cmd, mypypath)
@@ -366,6 +382,7 @@ def main() -> NoReturn:
         action="store_true",
         help="skip the test if the current platform is not specified in METADATA.toml/tool.stubtest.platforms",
     )
+    parser.add_argument("--keep-tmp-dir", action="store_true", help="keep the temporary virtualenv")
     parser.add_argument("dists", metavar="DISTRIBUTION", type=str, nargs=argparse.ZERO_OR_MORE)
     args = parser.parse_args()
 
@@ -379,7 +396,13 @@ def main() -> NoReturn:
     for i, dist in enumerate(dists):
         if i % args.num_shards != args.shard_index:
             continue
-        if not run_stubtest(dist, parser=parser, verbose=args.verbose, specified_platforms_only=args.specified_platforms_only):
+        if not run_stubtest(
+            dist,
+            parser=parser,
+            verbose=args.verbose,
+            specified_platforms_only=args.specified_platforms_only,
+            keep_tmp_dir=args.keep_tmp_dir,
+        ):
             result = 1
     sys.exit(result)
 
diff --git a/tests/typecheck_typeshed.py b/tests/typecheck_typeshed.py
index 6248c0e67..70d1d8e23 100755
--- a/tests/typecheck_typeshed.py
+++ b/tests/typecheck_typeshed.py
@@ -60,8 +60,9 @@ def run_mypy_as_subprocess(directory: str, platform: str, version: str) -> Retur
         "--no-error-summary",
         "--enable-error-code",
         "ignore-without-code",
-        "--enable-error-code",
-        "possibly-undefined",
+        # https://github.com/python/mypy/issues/14309
+        # "--enable-error-code",
+        # "possibly-undefined",
         "--enable-error-code",
         "redundant-expr",
         "--enable-error-code",