From 74112dc18919084093d7aed3f26456ccf4598cf8 Mon Sep 17 00:00:00 2001
From: Avasam <samuel.06@hotmail.com>
Date: Sun, 19 Feb 2023 17:10:32 -0500
Subject: [PATCH] Add a `run_stubtest` flag to `scripts/runtests.py` (#9717)

Co-authored-by: Alex Waygood <Alex.Waygood@Gmail.com>
---
 scripts/runtests.py | 50 ++++++++++++++++++++++++++++-----------------
 tests/README.md     |  3 +++
 2 files changed, 34 insertions(+), 19 deletions(-)

diff --git a/scripts/runtests.py b/scripts/runtests.py
index b2e2010d0..44e7e71a1 100644
--- a/scripts/runtests.py
+++ b/scripts/runtests.py
@@ -1,6 +1,7 @@
 #!/usr/bin/env python3
 from __future__ import annotations
 
+import argparse
 import json
 import os
 import re
@@ -44,16 +45,28 @@ def _get_strict_params(stub_path: str) -> list[str]:
 
 
 def main() -> None:
-    try:
-        path = sys.argv[1]
-    except IndexError:
-        print("Missing path argument in format <folder>/<stub>", file=sys.stderr)
-        sys.exit(1)
-    assert os.path.exists(path), rf"Path {path} does not exist."
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--run-stubtest",
+        action="store_true",
+        help=(
+            "Run stubtest for the selected package(s). Running stubtest may download and execute arbitrary code from PyPI: "
+            "only use this option if you trust the package you are testing."
+        ),
+    )
+    parser.add_argument("path", help="Path of the stub to test in format <folder>/<stub>, from the root of the project.")
+    args = parser.parse_args()
+    path: str = args.path
+    run_stubtest: bool = args.run_stubtest
+
     path_tokens = Path(path).parts
-    assert len(path_tokens) == 2, "Path argument should be in format <folder>/<stub>."
+    if len(path_tokens) != 2:
+        parser.error("'path' argument should be in format <folder>/<stub>.")
     folder, stub = path_tokens
-    assert folder in {"stdlib", "stubs"}, "Only the 'stdlib' and 'stubs' folders are supported."
+    if folder not in {"stdlib", "stubs"}:
+        parser.error("Only the 'stdlib' and 'stubs' folders are supported.")
+    if not os.path.exists(path):
+        parser.error(rf"'path' {path} does not exist.")
     stubtest_result: subprocess.CompletedProcess[bytes] | None = None
     pytype_result: subprocess.CompletedProcess[bytes] | None = None
 
@@ -99,19 +112,18 @@ def main() -> None:
             print("\nRunning stubtest...")
             stubtest_result = subprocess.run([sys.executable, "tests/stubtest_stdlib.py", stub])
         else:
-            run_stubtest_query = (
-                f"\nRun stubtest for {stub!r} (Y/N)?\n\n"
-                "NOTE: Running third-party stubtest involves downloading and executing arbitrary code from PyPI.\n"
-                f"Only run stubtest if you trust the {stub!r} package.\n"
-            )
-            run_stubtest_answer = input(colored(run_stubtest_query, "yellow")).lower()
-            while run_stubtest_answer not in {"yes", "no", "y", "n"}:
-                run_stubtest_answer = input(colored("Invalid response; please try again.\n", "red")).lower()
-            if run_stubtest_answer in {"yes", "y"}:
-                print("\nRunning stubtest.")
+            if run_stubtest:
+                print("\nRunning stubtest...")
                 stubtest_result = subprocess.run([sys.executable, "tests/stubtest_third_party.py", stub])
             else:
-                print(colored(f"\nSkipping stubtest for {stub!r}...", "yellow"))
+                print(
+                    colored(
+                        f"\nSkipping stubtest for {stub!r}..."
+                        + "\nNOTE: Running third-party stubtest involves downloading and executing arbitrary code from PyPI."
+                        + f"\nOnly run stubtest if you trust the {stub!r} package.",
+                        "yellow",
+                    )
+                )
     else:
         print(colored("\nSkipping stubtest since mypy failed.", "yellow"))
 
diff --git a/tests/README.md b/tests/README.md
index eb696b846..3abf11d4f 100644
--- a/tests/README.md
+++ b/tests/README.md
@@ -46,6 +46,9 @@ be selected. A summary of the results will be printed to the terminal.
 You must provide a single argument which is a path to the stubs to test, like
 so: `stdlib/os` or `stubs/requests`.
 
+Run `python scripts/runtests.py --help` for information on the various configuration options
+for this script.
+
 ## mypy\_test.py
 
 Run using: