[lit] Factor out listdir logic shared by different test formats.

Summary:
The lit test formats use largely the same logic for discovering tests. There are
some superficial differences in the logic, which seem reasonable enough to
handle in a single routine.

At a high level, the common goal is "look for files that end with one of these
suffixes, and skip anything starting with a dot." The balance of the logic
specific to ShTest and GoogleTest collapses quite a bit, so that
getTestsInDirectory is only a couple of lines around a call to the new function.

Reviewers: zturner, MatzeB, modocache

Subscribers: sanjoy, llvm-commits

Differential Revision: https://reviews.llvm.org/D34855

llvm-svn: 306895

3 files changed, 58 insertions, 44 deletions

diff --git a/llvm/utils/lit/lit/formats/googletest.py b/llvm/utils/lit/lit/formats/googletest.py
index f4efe2cf004..b683f7c7db8 100644
--- a/llvm/utils/lit/lit/formats/googletest.py
+++ b/llvm/utils/lit/lit/formats/googletest.py
@@ -11,8 +11,8 @@ from .base import TestFormat
 kIsWindows = sys.platform in ['win32', 'cygwin']
 
 class GoogleTest(TestFormat):
-    def __init__(self, test_sub_dir, test_suffix):
-        self.test_sub_dir = os.path.normcase(str(test_sub_dir)).split(';')
+    def __init__(self, test_sub_dirs, test_suffix):
+        self.test_sub_dirs = os.path.normcase(str(test_sub_dirs)).split(';')
         self.test_suffix = str(test_suffix)
 
         # On Windows, assume tests will also end in '.exe'.
@@ -74,38 +74,19 @@ class GoogleTest(TestFormat):
             else:
                 yield ''.join(nested_tests) + ln
 
-    # Note: path_in_suite should not include the executable name.
-    def getTestsInExecutable(self, testSuite, path_in_suite, execpath,
-                             litConfig, localConfig):
-        if not execpath.endswith(self.test_suffix):
-            return
-        (dirname, basename) = os.path.split(execpath)
-        # Discover the tests in this executable.
-        for testname in self.getGTestTests(execpath, litConfig, localConfig):
-            testPath = path_in_suite + (basename, testname)
-            yield lit.Test.Test(testSuite, testPath, localConfig, file_path=execpath)
-
     def getTestsInDirectory(self, testSuite, path_in_suite,
                             litConfig, localConfig):
         source_path = testSuite.getSourcePath(path_in_suite)
-        for filename in os.listdir(source_path):
-            filepath = os.path.join(source_path, filename)
-            if os.path.isdir(filepath):
-                # Iterate over executables in a directory.
-                if not os.path.normcase(filename) in self.test_sub_dir:
-                    continue
-                dirpath_in_suite = path_in_suite + (filename, )
-                for subfilename in os.listdir(filepath):
-                    execpath = os.path.join(filepath, subfilename)
-                    for test in self.getTestsInExecutable(
-                            testSuite, dirpath_in_suite, execpath,
-                            litConfig, localConfig):
-                      yield test
-            elif ('.' in self.test_sub_dir):
-                for test in self.getTestsInExecutable(
-                        testSuite, path_in_suite, filepath,
-                        litConfig, localConfig):
-                    yield test
+        for subdir in self.test_sub_dirs:
+            for fn in lit.util.listdir_files(os.path.join(source_path, subdir),
+                                             suffixes={self.test_suffix}):
+                # Discover the tests in this executable.
+                execpath = os.path.join(source_path, subdir, fn)
+                testnames = self.getGTestTests(execpath, litConfig, localConfig)
+                for testname in testnames:
+                    testPath = path_in_suite + (subdir, fn, testname)
+                    yield lit.Test.Test(testSuite, testPath, localConfig,
+                                        file_path=execpath)
 
     def execute(self, test, litConfig):
         testPath,testName = os.path.split(test.getSourcePath())
diff --git a/llvm/utils/lit/lit/formats/shtest.py b/llvm/utils/lit/lit/formats/shtest.py
index ea424f0f0cc..01ecd192092 100644
--- a/llvm/utils/lit/lit/formats/shtest.py
+++ b/llvm/utils/lit/lit/formats/shtest.py
@@ -4,6 +4,7 @@ import os
 
 import lit.Test
 import lit.TestRunner
+import lit.util
 from .base import TestFormat
 
 class ShTest(TestFormat):
@@ -34,19 +35,12 @@ class ShTest(TestFormat):
     def getTestsInDirectory(self, testSuite, path_in_suite,
                             litConfig, localConfig):
         """Yields test files matching 'suffixes' from the localConfig."""
-        source_path = testSuite.getSourcePath(path_in_suite)
-        for filename in os.listdir(source_path):
-            # Ignore dot files and excluded tests.
-            if (filename.startswith('.') or
-                filename in localConfig.excludes):
-                continue
-
-            filepath = os.path.join(source_path, filename)
-            if not os.path.isdir(filepath):
-                base,ext = os.path.splitext(filename)
-                if ext in localConfig.suffixes:
-                    yield lit.Test.Test(testSuite, path_in_suite + (filename,),
-                                        localConfig)
+        file_matches = lit.util.listdir_files(
+            testSuite.getSourcePath(path_in_suite),
+            localConfig.suffixes, localConfig.excludes)
+        for filename in file_matches:
+            yield lit.Test.Test(testSuite, path_in_suite + (filename,),
+                                localConfig)
 
     def execute(self, test, litConfig):
         """Interprets and runs the given test file, and returns the result."""
diff --git a/llvm/utils/lit/lit/util.py b/llvm/utils/lit/lit/util.py
index bc1a1a4bab0..1819d4d1c34 100644
--- a/llvm/utils/lit/lit/util.py
+++ b/llvm/utils/lit/lit/util.py
@@ -101,6 +101,45 @@ def mkdir_p(path):
         if e.errno != errno.EEXIST:
             raise
 
+def listdir_files(dirname, suffixes=None, exclude_filenames=None):
+    """Yields files in a directory.
+
+    Filenames that are not excluded by rules below are yielded one at a time, as
+    basenames (i.e., without dirname).
+
+    Files starting with '.' are always skipped.
+
+    If 'suffixes' is not None, then only filenames ending with one of its
+    members will be yielded. These can be extensions, like '.exe', or strings,
+    like 'Test'. (It is a lexicographic check; so an empty sequence will yield
+    nothing, but a single empty string will yield all filenames.)
+
+    If 'exclude_filenames' is not None, then none of the file basenames in it
+    will be yielded.
+
+    If specified, the containers for 'suffixes' and 'exclude_filenames' must
+    support membership checking for strs.
+
+    Args:
+        dirname: a directory path.
+        suffixes: (optional) a sequence of strings (set, list, etc.).
+        exclude_filenames: (optional) a sequence of strings.
+
+    Yields:
+        Filenames as returned by os.listdir (generally, str).
+    """
+    if exclude_filenames is None:
+        exclude_filenames = set()
+    if suffixes is None:
+        suffixes = {''}
+    for filename in os.listdir(dirname):
+        if (os.path.isdir(os.path.join(dirname, filename)) or
+            filename.startswith('.') or
+            filename in exclude_filenames or
+            not any(filename.endswith(sfx) for sfx in suffixes)):
+            continue
+        yield filename
+
 def which(command, paths = None):
     """which(command, [paths]) - Look up the given command in the paths string
     (or the PATH environment variable, if unspecified)."""