[lit] Small cleanups in main.py

[yln]

Summary

• Extract separate function for running tests from main
• Push single-usage imports to point of usage
• Remove unnecessary sys.exit(0) calls

[rnk]

lgtm

llvm/utils/lit/lit/main.py

32

Unrelated, but I wonder if we should augment this logic to garbage collect old lit_tmp_ directories that are 24+ hours old. I routinely find lots of leaked lit_tmp_ directories because oftentimes the parent Python process is killed before it gets to the finally block below.

[yln]

llvm/utils/lit/lit/main.py

32

One of the improvements for lit that I want to explore is graceful shutdown on CTRL+C. If that isn't possible, we should implement your suggestion.

llvm/utils/lit/lit/main.py

 #!/usr/bin/env python
 
 """
 lit - LLVM Integrated Tester.
 
 See lit.pod for more information.
 """
 
 from __future__ import absolute_import
 import os
 import platform
-import random
 import sys
 import time
-import tempfile
-import shutil
-from xml.sax.saxutils import quoteattr
 
 import lit.cl_arguments
 import lit.discovery
 import lit.display
 import lit.LitConfig
 import lit.run
 import lit.Test
 import lit.util
 
 def main(builtinParameters = {}):
     # Create a temp directory inside the normal temp directory so that we can
     # try to avoid temporary test file leaks. The user can avoid this behavior
     # by setting LIT_PRESERVES_TMP in the environment, so they can easily use
     # their own temp directory to monitor temporary file leaks or handle them at
     # the buildbot level.
     lit_tmp = None
     if 'LIT_PRESERVES_TMP' not in os.environ:
+        import tempfile
         lit_tmp = tempfile.mkdtemp(prefix="lit_tmp_")
         os.environ.update({
                 'TMPDIR': lit_tmp,
                 'TMP': lit_tmp,
                 'TEMP': lit_tmp,
                 'TEMPDIR': lit_tmp,
                 })
     # FIXME: If Python does not exit cleanly, this directory will not be cleaned
     # up. We should consider writing the lit pid into the temp directory,
     # scanning for stale temp directories, and deleting temp directories whose
     # lit process has died.
     try:
         main_with_tmp(builtinParameters)
     finally:
         if lit_tmp:
             try:
+                import shutil
                 shutil.rmtree(lit_tmp)
             except:
                 # FIXME: Re-try after timeout on Windows.
                 pass
 
 def main_with_tmp(builtinParameters):
     opts = lit.cl_arguments.parse_args()
 
def main_with_tmp(builtinParameters):
 
     # Finally limit the number of tests, if desired.
     if opts.maxTests is not None:
         run.tests = run.tests[:opts.maxTests]
 
     # Don't create more workers than tests.
     opts.numWorkers = min(len(run.tests), opts.numWorkers)
 
-    increase_process_limit(litConfig, opts)
-
-    display = lit.display.create_display(opts, len(run.tests),
-                                         numTotalTests, opts.numWorkers)
-    def progress_callback(test):
-        display.update(test)
-        if opts.incremental:
-            update_incremental_cache(test)
-
-    startTime = time.time()
-    try:
-        run.execute_tests(progress_callback, opts.numWorkers, opts.maxTime)
-    except KeyboardInterrupt:
-        sys.exit(2)
-    testing_time = time.time() - startTime
-
-    display.finish()
+    testing_time = run_tests(run, litConfig, opts, numTotalTests)
 
     if not opts.quiet:
         print('Testing Time: %.2fs' % (testing_time,))
 
     # Write out the test data, if requested.
     if opts.output_path is not None:
         write_test_results(run, litConfig, testing_time, opts.output_path)
 
if litConfig.numErrors:
         sys.exit(2)
 
     # Warn about warnings.
     if litConfig.numWarnings:
         sys.stderr.write('\n%d warning(s) in tests.\n' % litConfig.numWarnings)
 
     if hasFailures:
         sys.exit(1)
-    sys.exit(0)
 
 
 def create_user_parameters(builtinParameters, opts):
     userParams = dict(builtinParameters)
     for entry in opts.userParameters:
         if '=' not in entry:
             name,val = entry,''
         else:
def print_suites_or_tests(run, opts):
     # Show the tests, if requested.
     if opts.showTests:
         print('-- Available Tests --')
         for ts,ts_tests in suitesAndTests:
             ts_tests.sort(key = lambda test: test.path_in_suite)
             for test in ts_tests:
                 print('  %s' % (test.getFullName(),))
 
-    # Exit.
-    sys.exit(0)
-
 def order_tests(run, opts):
     if opts.shuffle:
+        import random
         random.shuffle(run.tests)
     elif opts.incremental:
         run.tests.sort(key = by_mtime, reverse = True)
     else:
         run.tests.sort(key = lambda t: (not t.isEarlyTest(), t.getFullName()))
 
 def by_mtime(test):
     fname = test.getFilePath()
try:
 
         if max_procs_soft < desired_limit:
             resource.setrlimit(resource.RLIMIT_NPROC, (desired_limit, max_procs_hard))
             litConfig.note('raised the process limit from %d to %d' % \
                                (max_procs_soft, desired_limit))
     except:
         pass
 
+def run_tests(run, litConfig, opts, numTotalTests):
+    increase_process_limit(litConfig, opts)
+
+    display = lit.display.create_display(opts, len(run.tests),
+                                         numTotalTests, opts.numWorkers)
+    def progress_callback(test):
+        display.update(test)
+        if opts.incremental:
+            update_incremental_cache(test)
+
+    startTime = time.time()
+    try:
+        run.execute_tests(progress_callback, opts.numWorkers, opts.maxTime)
+    except KeyboardInterrupt:
+        sys.exit(2)
+    testing_time = time.time() - startTime
+
+    display.finish()
+    return testing_time
+
 def write_test_results(run, lit_config, testing_time, output_path):
     try:
         import json
     except ImportError:
         lit_config.fatal('test output unsupported with Python 2.5')
 
     # Construct the data we will write.
     data = {}
def write_test_results(run, lit_config, testing_time, output_path):
     f = open(output_path, 'w')
     try:
         json.dump(data, f, indent=2, sort_keys=True)
         f.write('\n')
     finally:
         f.close()
 
 def write_test_results_xunit(run, opts):
+    from xml.sax.saxutils import quoteattr
     # Collect the tests, indexed by test suite
     by_suite = {}
     for result_test in run.tests:
         suite = result_test.suite.config.name
         if suite not in by_suite:
             by_suite[suite] = {
                                 'passes'   : 0,
                                 'failures' : 0,