[llvm] [lit] Add type hints to formats (PR #201583)

via llvm-commits llvm-commits at lists.llvm.org
Thu Jun 4 06:19:59 PDT 2026


https://github.com/prasoon054 created https://github.com/llvm/llvm-project/pull/201583

Currently `lit`'s test formats (`formats/base.py`, `formats/shtest.py` and `formats/googletest.py`) didn't have type annotations. This PR adds them.

>From f0382921afe3371310817dbdb5ac5fd4bb825652 Mon Sep 17 00:00:00 2001
From: Prasoon Kumar <prasoonkumar054 at gmail.com>
Date: Thu, 4 Jun 2026 18:42:43 +0530
Subject: [PATCH] [lit] Add type hints to formats

Signed-off-by: Prasoon Kumar <prasoonkumar054 at gmail.com>
---
 llvm/utils/lit/lit/formats/base.py       | 41 +++++++++++++++++++---
 llvm/utils/lit/lit/formats/googletest.py | 43 +++++++++++++++++++-----
 llvm/utils/lit/lit/formats/shtest.py     | 16 +++++++--
 3 files changed, 83 insertions(+), 17 deletions(-)

diff --git a/llvm/utils/lit/lit/formats/base.py b/llvm/utils/lit/lit/formats/base.py
index 6e53613e7d1cb..a8315fd7fb479 100644
--- a/llvm/utils/lit/lit/formats/base.py
+++ b/llvm/utils/lit/lit/formats/base.py
@@ -1,11 +1,24 @@
+from __future__ import annotations
+
 import os
+from typing import TYPE_CHECKING, Iterator, Tuple, Union
 
 import lit.Test
 import lit.util
 
+if TYPE_CHECKING:
+    from lit.LitConfig import LitConfig
+    from lit.TestingConfig import TestingConfig
+
 
 class TestFormat:
-    def getTestsForPath(self, testSuite, path_in_suite, litConfig, localConfig):
+    def getTestsForPath(
+        self,
+        testSuite: lit.Test.TestSuite,
+        path_in_suite: Tuple[str, ...],
+        litConfig: LitConfig,
+        localConfig: TestingConfig,
+    ) -> Iterator[lit.Test.Test]:
         """
         Given the path to a test in the test suite, generates the Lit tests associated
         to that path. There can be zero, one or more tests. For example, some testing
@@ -17,11 +30,18 @@ def getTestsForPath(self, testSuite, path_in_suite, litConfig, localConfig):
         """
         yield lit.Test.Test(testSuite, path_in_suite, localConfig)
 
+
 ###
 
 
 class FileBasedTest(TestFormat):
-    def getTestsForPath(self, testSuite, path_in_suite, litConfig, localConfig):
+    def getTestsForPath(
+        self,
+        testSuite: lit.Test.TestSuite,
+        path_in_suite: Tuple[str, ...],
+        litConfig: LitConfig,
+        localConfig: TestingConfig,
+    ) -> Iterator[lit.Test.Test]:
         """
         Expand each path in a test suite to a Lit test using that path and assuming
         it is a file containing the test. File extensions excluded by the configuration
@@ -36,20 +56,31 @@ def getTestsForPath(self, testSuite, path_in_suite, litConfig, localConfig):
         if any(filename.endswith(suffix) for suffix in localConfig.suffixes):
             yield lit.Test.Test(testSuite, path_in_suite, localConfig)
 
-    def getTestsInDirectory(self, testSuite, path_in_suite, litConfig, localConfig):
+    def getTestsInDirectory(
+        self,
+        testSuite: lit.Test.TestSuite,
+        path_in_suite: Tuple[str, ...],
+        litConfig: LitConfig,
+        localConfig: TestingConfig,
+    ) -> Iterator[lit.Test.Test]:
         source_path = testSuite.getSourcePath(path_in_suite)
         for filename in os.listdir(source_path):
             filepath = os.path.join(source_path, filename)
             if not os.path.isdir(filepath):
-                for t in self.getTestsForPath(testSuite, path_in_suite + (filename,), litConfig, localConfig):
+                for t in self.getTestsForPath(
+                    testSuite, path_in_suite + (filename,), litConfig, localConfig
+                ):
                     yield t
 
 
 ###
 
+
 # Check exit code of a simple executable with no input
 class ExecutableTest(FileBasedTest):
-    def execute(self, test, litConfig):
+    def execute(
+        self, test: lit.Test.Test, litConfig: LitConfig
+    ) -> Union[lit.Test.ResultCode, Tuple[lit.Test.ResultCode, str]]:
         if test.config.unsupported:
             return lit.Test.UNSUPPORTED
 
diff --git a/llvm/utils/lit/lit/formats/googletest.py b/llvm/utils/lit/lit/formats/googletest.py
index e567cce541beb..cacd6ccccfcce 100644
--- a/llvm/utils/lit/lit/formats/googletest.py
+++ b/llvm/utils/lit/lit/formats/googletest.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import json
 import math
 import os
@@ -9,12 +11,23 @@
 import lit.TestRunner
 import lit.util
 from .base import TestFormat
+from typing import TYPE_CHECKING, Dict, Iterator, List, Optional, Tuple, Union
+
+if TYPE_CHECKING:
+    from lit.LitConfig import LitConfig
+    from lit.TestingConfig import TestingConfig
 
 kIsWindows = sys.platform in ["win32", "cygwin"]
 
 
 class GoogleTest(TestFormat):
-    def __init__(self, test_sub_dirs, test_suffix, run_under=[], test_prefix=None):
+    def __init__(
+        self,
+        test_sub_dirs: str,
+        test_suffix: str,
+        run_under: Union[str, List[str]] = [],
+        test_prefix: Optional[str] = None,
+    ) -> None:
         self.seen_executables = set()
         self.test_sub_dirs = str(test_sub_dirs).split(";")
 
@@ -28,7 +41,9 @@ def __init__(self, test_sub_dirs, test_suffix, run_under=[], test_prefix=None):
         self.test_prefixes = {test_prefix} if test_prefix else None
         self.run_under = run_under
 
-    def get_num_tests(self, path, litConfig, localConfig):
+    def get_num_tests(
+        self, path: str, litConfig: LitConfig, localConfig: TestingConfig
+    ) -> Optional[int]:
         list_test_cmd = self.prepareCmd(
             [path, "--gtest_list_tests", "--gtest_filter=-*DISABLED_*"]
         )
@@ -47,7 +62,13 @@ def get_num_tests(self, path, litConfig, localConfig):
             )
         )
 
-    def getTestsInDirectory(self, testSuite, path_in_suite, litConfig, localConfig):
+    def getTestsInDirectory(
+        self,
+        testSuite: lit.Test.TestSuite,
+        path_in_suite: Tuple[str, ...],
+        litConfig: LitConfig,
+        localConfig: TestingConfig,
+    ) -> Iterator[lit.Test.Test]:
         init_shard_size = 512  # number of tests in a shard
         core_count = lit.util.usable_core_count()
         source_path = testSuite.getSourcePath(path_in_suite)
@@ -141,7 +162,9 @@ def getTestsInDirectory(self, testSuite, path_in_suite, litConfig, localConfig):
                         testSuite, testPath, localConfig, file_path=execpath
                     )
 
-    def execute(self, test, litConfig):
+    def execute(
+        self, test: lit.Test.Test, litConfig: LitConfig
+    ) -> Tuple[lit.Test.ResultCode, str]:
         if test.gtest_json_file is None:
             return lit.Test.FAIL, ""
 
@@ -181,7 +204,7 @@ def execute(self, test, litConfig):
         if litConfig.noExecute:
             return lit.Test.PASS, ""
 
-        def get_shard_header(shard_env):
+        def get_shard_header(shard_env: Dict[str, str]) -> str:
             shard_envs = " ".join([k + "=" + v for k, v in shard_env.items()])
             return f"Script(shard):\n--\n%s %s\n--\n" % (shard_envs, " ".join(cmd))
 
@@ -210,7 +233,7 @@ def get_shard_header(shard_env):
         if exitCode == 0:
             return lit.Test.PASS, ""
 
-        def get_test_stdout(test_name):
+        def get_test_stdout(test_name: str) -> str:
             res = []
             header = f"[ RUN      ] " + test_name
             footer = f"[  FAILED  ] " + test_name
@@ -263,7 +286,7 @@ def get_test_stdout(test_name):
 
         return lit.Test.FAIL, output
 
-    def prepareCmd(self, cmd):
+    def prepareCmd(self, cmd: List[str]) -> List[str]:
         """Insert interpreter if needed.
 
         It inserts the python exe into the command if cmd[0] ends in .py or caller
@@ -282,8 +305,10 @@ def prepareCmd(self, cmd):
         return cmd
 
     @staticmethod
-    def post_process_shard_results(selected_tests, discovered_tests):
-        def remove_gtest(tests):
+    def post_process_shard_results(
+        selected_tests: List[lit.Test.Test], discovered_tests: List[lit.Test.Test]
+    ) -> Tuple[List[lit.Test.Test], List[lit.Test.Test]]:
+        def remove_gtest(tests: List[lit.Test.Test]) -> List[lit.Test.Test]:
             return [t for t in tests if t.gtest_json_file is None]
 
         discovered_tests = remove_gtest(discovered_tests)
diff --git a/llvm/utils/lit/lit/formats/shtest.py b/llvm/utils/lit/lit/formats/shtest.py
index 5529b75f51d4f..b17331bf15664 100644
--- a/llvm/utils/lit/lit/formats/shtest.py
+++ b/llvm/utils/lit/lit/formats/shtest.py
@@ -1,7 +1,14 @@
+from __future__ import annotations
+
 import lit.TestRunner
 import lit.util
 
 from .base import FileBasedTest
+from typing import TYPE_CHECKING, List, Tuple
+
+if TYPE_CHECKING:
+    from lit.Test import Test, Result
+    from lit.LitConfig import LitConfig
 
 
 class ShTest(FileBasedTest):
@@ -17,13 +24,16 @@ class ShTest(FileBasedTest):
     """
 
     def __init__(
-        self, execute_external=False, extra_substitutions=[], preamble_commands=[]
-    ):
+        self,
+        execute_external: bool = False,
+        extra_substitutions: List[Tuple[str, str]] = [],
+        preamble_commands: List[str] = [],
+    ) -> None:
         self.execute_external = execute_external
         self.extra_substitutions = extra_substitutions
         self.preamble_commands = preamble_commands
 
-    def execute(self, test, litConfig):
+    def execute(self, test: Test, litConfig: LitConfig) -> Result:
         return lit.TestRunner.executeShTest(
             test,
             litConfig,



More information about the llvm-commits mailing list