[mypyc] Support _benchmark test suffix for optimized run tests (#21047)

If a mypyc run test has `_benchmark` in the test name as a suffix (not
necessarily the final suffix), the test case will be compiled with
optimizations enabled, and this also applies to `librt` if it needs to
be built. The idea is that this can be used to easily run `librt`
benchmarks by using temporary throwaway tests.

Tested by manually running some test cases with and without a
`_benchmark` suffix.

Author: JukkaL

mypyc/test/librt_cache.py

@@ -33,14 +33,15 @@
 from mypyc.common import RUNTIME_C_FILES
 
 
-def _librt_build_hash(experimental: bool) -> str:
+def _librt_build_hash(experimental: bool, opt_level: str) -> str:
     """Compute hash for librt build, including sources and build environment."""
     # Import lazily to ensure mypyc.build has ensured that distutils is correctly set up
     from distutils import ccompiler
 
     h = hashlib.sha256()
     # Include experimental flag
     h.update(b"exp" if experimental else b"noexp")
+    h.update(f"opt={opt_level}".encode())
     # Include full Python version string (includes git hash for dev builds)
     h.update(sys.version.encode())
     # Include debug build status (gettotalrefcount only exists in debug builds)
@@ -72,16 +73,16 @@ def _librt_build_hash(experimental: bool) -> str:
     return h.hexdigest()[:16]
 
 
-def _generate_setup_py(build_dir: str, experimental: bool) -> str:
+def _generate_setup_py(build_dir: str, experimental: bool, opt_level: str) -> str:
     """Generate setup.py content for building librt directly.
 
     We inline LIBRT_MODULES/RUNTIME_C_FILES/include_dir/cflags values to avoid
     importing mypyc.build, which recursively imports lots of things.
     """
     lib_rt_dir = include_dir()
 
-    # Get compiler flags using the shared helper (with -O0 for faster builds)
-    cflags = get_cflags(opt_level="0", experimental_features=experimental)
+    # Get compiler flags using the shared helper
+    cflags = get_cflags(opt_level=opt_level, experimental_features=experimental)
 
     # Serialize values to inline in generated setup.py
     librt_modules_repr = repr(
@@ -135,7 +136,7 @@ def write_file(path, contents):
 """
 
 
-def get_librt_path(experimental: bool = True) -> str:
+def get_librt_path(experimental: bool = True, opt_level: str = "0") -> str:
     """Get path to librt built from the repository, building and caching if necessary.
 
     Uses build/librt-cache/ under the repo root (gitignored). The cache is
@@ -146,14 +147,15 @@ def get_librt_path(experimental: bool = True) -> str:
 
     Args:
         experimental: Whether to enable experimental features.
+        opt_level: Optimization level ("0".."3") used when building librt.
 
     Returns:
         Path to directory containing built librt modules.
     """
     # Use build/librt-cache/ under the repo root (gitignored)
     repo_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
     cache_root = os.path.join(repo_root, "build", "librt-cache")
-    build_hash = _librt_build_hash(experimental)
+    build_hash = _librt_build_hash(experimental, opt_level)
     build_dir = os.path.join(cache_root, f"librt-{build_hash}")
     lock_file = os.path.join(cache_root, f"librt-{build_hash}.lock")
     marker = os.path.join(build_dir, ".complete")
@@ -186,7 +188,7 @@ def get_librt_path(experimental: bool = True) -> str:
         # Write setup.py
         setup_py = os.path.join(build_dir, "setup.py")
         with open(setup_py, "w") as f:
-            f.write(_generate_setup_py(build_dir, experimental))
+            f.write(_generate_setup_py(build_dir, experimental, opt_level))
 
         # Build (parallel builds don't work well because multiple extensions
         # share the same runtime C files, causing race conditions)
@@ -207,7 +209,7 @@ def get_librt_path(experimental: bool = True) -> str:
 
 
 def run_with_librt(
-    file_path: str, experimental: bool = True, check: bool = True
+    file_path: str, experimental: bool = True, check: bool = True, opt_level: str = "0"
 ) -> subprocess.CompletedProcess[str]:
     """Run a Python file in a subprocess with built librt available.
 
@@ -218,11 +220,12 @@ def run_with_librt(
         file_path: Path to Python file to execute.
         experimental: Whether to use experimental features.
         check: If True, raise CalledProcessError on non-zero exit.
+        opt_level: Optimization level ("0".."3") used when building librt.
 
     Returns:
         CompletedProcess with stdout, stderr, and returncode.
     """
-    librt_path = get_librt_path(experimental)
+    librt_path = get_librt_path(experimental, opt_level=opt_level)
     # Prepend librt path to PYTHONPATH
     env = os.environ.copy()
     existing = env.get("PYTHONPATH", "")

mypyc/test/test_run.py

@@ -36,6 +36,7 @@
     MypycDataSuite,
     assert_test_output,
     fudge_dir_mtimes,
+    has_test_name_tag,
     show_c,
     use_custom_builtins,
 )
@@ -210,7 +211,10 @@ def run_case_inner(self, testcase: DataDrivenTestCase) -> None:
             self.run_case_step(testcase, step)
 
     def run_case_step(self, testcase: DataDrivenTestCase, incremental_step: int) -> None:
-        bench = testcase.config.getoption("--bench", False) and "Benchmark" in testcase.name
+        benchmark_build = has_test_name_tag(testcase.name, "benchmark")
+        bench = testcase.config.getoption("--bench", False) and (
+            benchmark_build or "Benchmark" in testcase.name
+        )
 
         options = Options()
         options.use_builtins_fixtures = True
@@ -262,10 +266,10 @@ def run_case_step(self, testcase: DataDrivenTestCase, incremental_step: int) ->
 
         # Use _librt_internal to test mypy-specific parts of librt (they have
         # some special-casing in mypyc), for everything else use _librt suffix.
-        librt_internal = testcase.name.endswith("_librt_internal")
-        librt = testcase.name.endswith("_librt") or "_librt_" in testcase.name
+        librt_internal = has_test_name_tag(testcase.name, "librt_internal")
+        librt = has_test_name_tag(testcase.name, "librt")
         # Enable experimental features (local librt build also includes experimental features)
-        experimental_features = testcase.name.endswith("_experimental")
+        experimental_features = has_test_name_tag(testcase.name, "experimental")
         try:
             compiler_options = CompilerOptions(
                 multi_file=self.multi_file,
@@ -302,7 +306,7 @@ def run_case_step(self, testcase: DataDrivenTestCase, incremental_step: int) ->
         if incremental_step == 1:
             check_serialization_roundtrip(ir)
 
-        opt_level = int(os.environ.get("MYPYC_OPT_LEVEL", 0))
+        opt_level = 3 if benchmark_build else int(os.environ.get("MYPYC_OPT_LEVEL", 0))
 
         setup_file = os.path.abspath(os.path.join(WORKDIR, "setup.py"))
         # We pass the C file information to the build script via setup.py unfortunately
@@ -322,7 +326,7 @@ def run_case_step(self, testcase: DataDrivenTestCase, incremental_step: int) ->
 
         if librt:
             # Use cached pre-built librt instead of rebuilding for each test
-            cached_librt = get_librt_path(experimental_features)
+            cached_librt = get_librt_path(experimental_features, opt_level=str(opt_level))
             shutil.copytree(os.path.join(cached_librt, "librt"), "librt")
 
         if not run_setup(setup_file, ["build_ext", "--inplace"]):

mypyc/test/testutil.py

@@ -286,6 +286,15 @@ def infer_ir_build_options_from_test_name(name: str) -> CompilerOptions | None:
         options.python_version = options.capi_version
     elif "_py" in name or "_Python" in name:
         assert False, f"Invalid _py* suffix (should be _pythonX_Y): {name}"
-    if re.search("_experimental(_|$)", name):
+    if has_test_name_tag(name, "experimental"):
         options.experimental_features = True
     return options
+
+
+def has_test_name_tag(name: str, tag: str) -> bool:
+    """Check if a test case name contains a tag token like ``_experimental``.
+
+    A tag matches if it appears as a full underscore-delimited token:
+    ``foo_tag_bar`` or ``foo_tag``.
+    """
+    return re.search(rf"(?:^|_){re.escape(tag)}(?:_|$)", name) is not None