Index: llvm/docs/TestingGuide.rst
===================================================================
--- llvm/docs/TestingGuide.rst
+++ llvm/docs/TestingGuide.rst
@@ -219,9 +219,16 @@
 only directories does not need the ``lit.local.cfg`` file. Read the :doc:`Lit
 documentation <CommandGuide/lit>` for more information.
 
-Each test file must contain lines starting with "RUN:" that tell :program:`lit`
-how to run it. If there are no RUN lines, :program:`lit` will issue an error
-while running a test.
+Each test file must contain lines with "RUN:" or "PYTHON:" that tell
+:program:`lit` how to run it, as described in the following sections. If there
+are no such lines, :program:`lit` will issue an error while running a test.
+
+Put related tests into a single file rather than having a separate file per
+test. Check if there are files already covering your feature and consider
+adding your code there instead of creating a new file.
+
+RUN lines
+---------
 
 RUN lines are specified in the comments of the test program using the
 keyword ``RUN`` followed by a colon, and lastly the command (pipeline)
@@ -271,9 +278,205 @@
 the :doc:`FileCheck tool <CommandGuide/FileCheck>`. *[The usage of grep in RUN
 lines is deprecated - please do not send or commit patches that use it.]*
 
-Put related tests into a single file rather than having a separate file per
-test. Check if there are files already covering your feature and consider
-adding your code there instead of creating a new file.
+PYTHON lines
+------------
+
+While RUN lines are great for specifying LLVM tool command lines, its simple
+shell-like functionality can prove insufficient for some testing logic. For such
+cases, PYTHON lines offer the full power of python scripting, including loops,
+functions, and modules, as illustrated in the examples below.
+
+**lit.run(cmd):**
+
+In order to move a RUN line into, for example, a python control structure like a
+loop, you must convert it to a ``lit.run(cmd)`` call in a PYTHON line. For
+example:
+
+.. code-block:: llvm
+
+    ; PYTHON: for triple in ("x86_64-apple-darwin10.6.0", "x86_64-unknown-linux-gnu"):
+    ; PYTHON:     lit.run(f"""
+    ; PYTHON:         %clang_cc1 -verify -fopenmp -fopenmp-version=51
+    ; PYTHON:                    -triple {triple} -emit-llvm -o - %s |
+    ; PYTHON:           FileCheck %s
+    ; PYTHON:     """)
+
+``cmd`` in ``PYTHON: lit.run(cmd)`` behaves very much like it does in ``RUN:
+cmd``. That is, :program:`lit` performs the same substitutions, parses the same
+shell-like syntax, and then executes the resulting commands. Of course, a
+difference is that python string syntax must be followed in PYTHON lines. For
+example, in a python f-string, a :program:`lit` substitution like ``%{foo}``
+must be escaped as ``%{{foo}}``.
+
+**Python functions:**
+
+The ability to write loops in PYTHON lines makes it easy to fall into the trap
+of creating long loop nests to achieve test coverage for all points in a large
+problem space. Especially given that :program:`lit` serially executes all
+commands in a single test file (:program:`lit` parallelism is only across test
+files), this approach can produce very slow tests. Unless full coverage is
+critical, it might be better to specify only a cross-section of such a problem
+space. Python functions can be useful for this purpose. For example, we might
+extend the above example with SIMD testing but choose to do so for only one
+triple:
+
+.. code-block:: llvm
+
+    ; PYTHON: def check(cflags, fcprefix):
+    ; PYTHON:     lit.run(f"""
+    ; PYTHON:         %clang_cc1 -verify -fopenmp -fopenmp-version=51
+    ; PYTHON:                    {cflags} -emit-llvm -o - %s |
+    ; PYTHON:           FileCheck -check-prefix={fcprefix} %s
+    ; PYTHON:     """)
+    ; PYTHON:
+    ; PYTHON: check("-triple x86_64-apple-darwin10.6.0 -fopenmp-simd", "SIMD")
+    ; PYTHON: check("-triple x86_64-apple-darwin10.6.0", "NO-SIMD")
+    ; PYTHON: check("-triple x86_64-unknown-linux-gnu", "NO-SIMD")
+
+**Python symbol scope:**
+
+The above example defines a new python symbol, the function ``check``. As in a
+python script, symbols defined in a PYTHON line are visible in all later PYTHON
+lines in the same test file. They are not visible in other test files, which is
+especially important given that :program:`lit` test file execution order is
+usually non-deterministic.
+
+**Importing python modules:**
+
+You might decide that some python symbols need to be shared across a set of test
+files. As usual in python, you can define such symbols in a python module (but
+see `config.prologue`_ below for an alternative approach if you need to share
+them across an entire test directory). For example, if the module file
+``exampleModule.py`` is located in the same directory as your test files, you
+can import it as follows:
+
+.. code-block:: llvm
+
+    ; PYTHON: import sys
+    ; PYTHON: sys.path.append(lit.expand("%S"))
+    ; PYTHON: import exampleModule
+
+**lit.expand(text):**
+
+``lit.expand(text)`` performs the same substitutions that :program:`lit`
+performs in RUN lines and in `lit.run(cmd)`. In the above example, it retrieves
+the value of :program:`lit`'s ``%S`` substitution, which is the test case's
+source directory, so that ``exampleModule.py`` can be located.
+
+**Exposing the lit object to an imported module:**
+
+As usual in python, an imported module does not automatically have access to
+symbols in the scope of the importer. In the above example, ``exampleModule.py``
+does not have access to the ``lit`` object. However, it needs access if it
+defines the ``check`` function above.
+
+One solution is to declare ``lit`` within the module and then assign it in a
+PYTHON line. For example, the contents of ``exampleModule.py`` would be as
+follows:
+
+.. code-block:: python
+
+    lit = None
+
+    def check(cflags, fcprefix):
+        lit.run(f"""
+            %clang_cc1 -verify -fopenmp -fopenmp-version=51
+                       {cflags} -emit-llvm -o - %s |
+              FileCheck -check-prefix={fcprefix} %s
+        """)
+
+Then your test's PYTHON lines would import ``exampleModule.py`` as follows:
+
+.. code-block:: llvm
+
+    ; PYTHON: import sys
+    ; PYTHON: sys.path.append(lit.expand("%S"))
+    ; PYTHON: import exampleModule
+    ; PYTHON: exampleModule.lit = lit
+
+.. _config.prologue:
+
+**config.prologue:**
+
+If ``exampleModule.py`` is needed in every test file in a test directory, the
+above import code might be cumbersome to repeat and maintain as a prologue in
+them all, especially if it evolves over time to include other modules or to
+perform some additional initialization.
+
+To avoid this problem, you can set ``config.prologue`` in a lit configuration
+file. Its value must be the name of a python file, whose contents :program:`lit`
+handles as if they appear in PYTHON lines at the start of every test file. Thus,
+any python symbols it defines are seen throughout every test file, and it has
+access to the ``lit`` object. For example, you might decide to rename
+``exampleModule.py`` to something like ``lit.prologue.py``, set
+``config.prologue`` to that file, and thus avoid writing the code for importing
+a module and exposing the ``lit`` object.
+
+**lit.has(feature):**
+
+To determine if ``feature`` appears in ``config.available_features`` (see
+`Constraining test execution`_ and `%if` under `Substitutions`_), you can call
+`lit.has(feature)`.
+
+**Mixing PYTHON lines and other directives:**
+
+It is fine to mix RUN lines and PYTHON lines in the same test file.
+:program:`lit` executes them in the order in which they appear. A possible way
+for them to interact is via temporary files (e.g., see ``%t`` in
+`Substitutions`_).
+
+It is also fine to mix PYTHON lines with DEFINE and REDEFINE lines, which assign
+:program:`lit` substitutions (see `Test-specific substitutions`_).
+`lit.expand(text)` and `lit.run(cmd)` will see the latest substitution values
+just as RUN lines do.
+
+The ability to mix PYTHON lines with other directives is helpful for at least
+two reasons. First, it facilitates extending existing test files with PYTHON
+lines without having to replace all RUN, DEFINE, and REDEFINE lines. Second,
+even for new tests, when a simple ``RUN: cmd`` would be sufficient, it is
+probably best not to complicate it with python syntax by converting it to
+``PYTHON: lit.run(cmd)``.
+
+**Python blocks and indentation:**
+
+:program:`lit` parses and executes all consecutive PYTHON lines as a single
+PYTHON block. For :program:`lit`, consecutive just means there are no other
+:program:`lit` directives (e.g., RUN lines) in between, but :program:`lit`
+discards other lines as inert text. The python code in a PYTHON block must be
+syntactically complete. For example, a PYTHON block cannot end with ``PYTHON: if
+cond:`` no matter what appears at the start of the next PYTHON block.
+
+Any whitespace immediately after ``PYTHON:`` on the initial line of a PYTHON
+block must also appear immediately after ``PYTHON:`` on all non-initial lines in
+the block. :program:`lit` trims that whitespace from every line before passing
+it to the python interpreter. Thus, non-initial lines can contain additional
+indentation according to the rules of python syntax.
+
+**Calling python from RUN lines:**
+
+In some tests, you might find techniques for calling python code from
+a RUN line rather than from a PYTHON line.  In particular, a test
+might call the python interpreter (usually specified via a ``%python``
+or ``%{python}`` :program:`lit` substitution) on an external python
+file.  A test might use the ``split-file`` utility (see `Extra files`_
+below) to embed that python file directly into the test file if the
+python code is sufficiently concise and does not need to be reused
+across multiple tests.
+
+A disadvantage of such techniques is that they do not have access to
+the ``lit`` object exposed to PYTHON lines.  For example, to move a
+shell command from a RUN line into a control structure like a loop or
+function without using ``lit.run(cmd)`` as in the examples above, the
+test author can instead pass the shell command to python's own
+``subprocess.run`` function.  To make any required :program:`lit`
+substitutions available, they can pass them as command-line arguments
+on the python interpreter RUN line, keeping in mind that substitutions
+are sometimes complex and that proper shell quoting must be applied.
+It is the test author's job to decide how to capture stdout, stderr,
+exit status, etc. from the shell command and present them for
+debugging when the test fails.  In contrast, ``lit.run(cmd)`` called
+from a PYTHON line handles all these issues already in a manner that
+is consistent with :program:`lit`'s handling of RUN lines.
 
 Generating assertions in regression tests
 -----------------------------------------
@@ -788,6 +991,9 @@
   ``DEFINE:`` and ``REDEFINE:`` directives, described in detail below.  So that
   they have no effect on other test files, these directives modify a copy of the
   substitution list that is produced by lit configuration files.
+- While simple substitution definitions can be useful for avoiding repeated
+  code, complex substitution relationships can be hard to understand and
+  maintain.  Consider using python functions instead (see `PYTHON lines`_).
 
 For example, the following directives can be inserted into a test file to define
 ``%{cflags}`` and ``%{fcflags}`` substitutions with empty initial values, which
Index: llvm/utils/lit/lit/TestRunner.py
===================================================================
--- llvm/utils/lit/lit/TestRunner.py
+++ llvm/utils/lit/lit/TestRunner.py
@@ -1,5 +1,7 @@
 from __future__ import absolute_import
+import contextlib
 import errno
+import inspect
 import io
 import itertools
 import getopt
@@ -12,6 +14,7 @@
 import shutil
 import tempfile
 import threading
+import traceback
 
 import io
 
@@ -45,6 +48,18 @@
         super().__init__(message)
 
 
+class ScriptFail(Exception):
+    """
+    A script failed, but it might by worth retrying.  Any failure message has
+    already been emitted on stdout or stderr.
+    """
+
+    def __init__(self, exitCode, timeoutInfo):
+        self.exitCode = exitCode
+        self.timeoutInfo = timeoutInfo
+        self.status = Test.FAIL if timeoutInfo is None else Test.TIMEOUT
+
+
 kIsWindows = platform.system() == "Windows"
 
 # Don't use close_fds on Windows.
@@ -66,6 +81,8 @@
 #
 # COMMAND that follows %dbg(ARG) is also captured. COMMAND can be
 # empty as a result of conditinal substitution.
+#
+# An empty '%dbg()' means sufficient debug info was already printed.
 kPdbgRegex = "%dbg\\(([^)'\"]*)\\)((?:.|\\n)*)"
 
 
@@ -1033,8 +1050,10 @@
     cmds = []
     for i, ln in enumerate(commands):
         # Within lit, we try to always add '%dbg(...)' to command lines in order
-        # to maximize debuggability.  However, custom lit test formats might not
-        # always add it, so add a generic debug message in that case.
+        # to maximize debuggability, or we add '%dbg()' to indicate a debug
+        # message was already printed in the execution trace.  However, custom
+        # lit test formats might not always add it, so add a generic debug
+        # message in that case.
         match = re.fullmatch(kPdbgRegex, ln)
         if match:
             dbg = match.group(1)
@@ -1043,11 +1062,15 @@
             dbg = "command line"
             command = ln
         if debug:
-            ln = f"@echo '# {dbg}' "
+            ln = ""
             if command:
-                ln += f"&& @echo {shlex.quote(command.lstrip())} && {command}"
+                if dbg:  # non-empty %dbg(...)
+                    ln += f"@echo '# {dbg}' && "
+                ln += f"@echo {shlex.quote(command.lstrip())} && {command}"
             else:
-                ln += "has no command after substitutions"
+                if not dbg:  # empty %dbg()
+                    dbg = "command line"
+                ln += f"@echo '# {dbg} has no command after substitutions'"
         else:
             ln = command
         try:
@@ -1458,7 +1481,9 @@
 
         '\' is documented as indicating a line continuation even if whitespace
         separates it from the newline.  It looks like a line continuation, and
-        it would be confusing if it didn't behave as one.
+        it would be confusing if it didn't behave as one.  (However, for python
+        directives, we just let the python compiler decide how to handle this
+        case.)
         """
         assert False, "expected method to be called on derived class"
 
@@ -1621,6 +1646,102 @@
         substitutions[existing[0]] = (self.name, value_repl)
 
 
+class PythonDirective(ExpandableScriptDirective):
+    """
+    A lit directive taking a python statement.  For example,
+    'PYTHON: lit.run(cmd)'.
+
+    indent: The indentation found on the first line and expected on continuation
+        lines.  It must be stripped for python to compile the code.
+    body: The uncompiled code accumulated so far from the directive and its
+        continuation lines.
+    """
+
+    def __init__(self, start_line_number, end_line_number, keyword, line):
+        super().__init__(start_line_number, end_line_number, keyword)
+        # Add blank lines so python diagnostics produce correct line numbers.
+        self.body = (start_line_number - 1) * "\n"
+        # Determine the indentation and remove it so that it's valid python.
+        line_stripped = line.lstrip()
+        self.indent = line[0 : len(line) - len(line_stripped)]
+        self.body += line_stripped + "\n"
+
+    def add_continuation(self, line_number, keyword, line):
+        if keyword != self.keyword:
+            return False
+        # Add blank lines so python diagnostics produce correct line numbers.
+        assert (
+            self.end_line_number < line_number
+        ), "expected monotonically increasing line number"
+        self.body += (line_number - self.end_line_number - 1) * "\n"
+        self.end_line_number = line_number
+        # Remove common leading indentation so it's valid python.
+        if not line.startswith(self.indent):
+            raise ValueError(
+                f"'{self.keyword}' directive continuation line has indentation "
+                f"that is inconsistent with previous lines"
+            )
+        self.body += line[len(self.indent) :] + "\n"
+        return True
+
+    def needs_continuation(self):
+        # Unlike lit directives that always mark a continued line with a
+        # trailing '\\', sometimes you cannot tell whether a python directive is
+        # continued until you see the indentation of the next python directive.
+        # In general, lit shouldn't reimplement python's line continuation and
+        # block syntax handling.  Thus, lit just concatenates all consecutive
+        # python directives and compiles them together even if they could be
+        # compiled separately.
+        return False
+
+    @staticmethod
+    def executePython(what, filename, pythonCode, python_dict, log):
+        """
+        Execute python code passed as an argument.
+        """
+        python_dict["__file__"] = filename = os.path.abspath(filename)
+        try:
+            c = compile(pythonCode, filename, "exec")
+        except Exception as e:
+            # Report the exception without this stack frame.  Bothering the lit
+            # user with lit internals is just a distraction.
+            tb = e.__traceback__.tb_next
+            traceText = "".join(traceback.format_exception(None, e, tb))
+            raise ScriptFatal(f"error compiling {what}:\n{traceText}")
+        log(f"# executed {what}\n")
+        outIO = StringIO()
+        errIO = StringIO()
+        traceText = ""
+        try:
+            with contextlib.redirect_stdout(outIO), contextlib.redirect_stderr(errIO):
+                exec(c, python_dict)
+        except ScriptFail:
+            raise
+        except Exception as e:
+            # Report the exception without this stack frame.  Bothering the lit
+            # user with lit internals is just a distraction.
+            tb = e.__traceback__.tb_next
+            traceText = "".join(traceback.format_exception(None, e, tb))
+            raise ScriptFail(1, None)
+        finally:
+            outText = outIO.getvalue()
+            errText = errIO.getvalue() + traceText
+            log(formatOutput(f"stdout from {what}", outText))
+            log(formatOutput(f"stderr from {what}", errText))
+
+    def execute(self, filename, python_dict, log):
+        """
+        Execute the python code in this directive.
+        """
+        self.executePython(
+            f"'{self.keyword}' directive {self.get_location()}",
+            filename,
+            self.body,
+            python_dict,
+            log,
+        )
+
+
 def applySubstitutions(script, substitutions, conditions={}, recursion_limit=None):
     """
     Apply substitutions to the script.  Allow full regular expression syntax.
@@ -1775,7 +1896,7 @@
             if isinstance(directive, CommandDirective):
                 line = directive.command
             else:
-                # Can come from preamble_commands.
+                # Can come from preamble_commands or from lit.expand.
                 assert isinstance(directive, str)
                 line = directive
             output.append(unescapePercents(process(line)))
@@ -1799,6 +1920,7 @@
         'DEFINE: %{name}=value'
     REDEFINE: A keyword taking a lit substitution redefinition. Ex
         'REDEFINE: %{name}=value'
+    PYTHON: A keyword taking a python statement. Ex 'PYTHON: lit.run(cmd)'
     """
 
     TAG = 0
@@ -1809,6 +1931,7 @@
     CUSTOM = 5
     DEFINE = 6
     REDEFINE = 7
+    PYTHON = 8
 
     @staticmethod
     def allowedKeywordSuffixes(value):
@@ -1821,6 +1944,7 @@
             ParserKind.CUSTOM: [":", "."],
             ParserKind.DEFINE: [":"],
             ParserKind.REDEFINE: [":"],
+            ParserKind.PYTHON: [":"],
         }[value]
 
     @staticmethod
@@ -1834,6 +1958,7 @@
             ParserKind.CUSTOM: "CUSTOM",
             ParserKind.DEFINE: "DEFINE",
             ParserKind.REDEFINE: "REDEFINE",
+            ParserKind.PYTHON: "PYTHON",
         }[value]
 
 
@@ -1895,6 +2020,10 @@
             self.parser = lambda line_number, line, output: self._handleSubst(
                 line_number, line, output, self.keyword, new_subst=False
             )
+        elif kind == ParserKind.PYTHON:
+            self.parser = lambda line_number, line, output: self._handlePython(
+                line_number, line, output, self.keyword
+            )
         else:
             raise ValueError("Unknown kind '%s'" % kind)
 
@@ -1997,6 +2126,16 @@
         )
         return output
 
+    @classmethod
+    def _handlePython(cls, line_number, line, output, keyword):
+        """A parser for PYTHON type keywords"""
+        if output and output[-1].add_continuation(line_number, keyword, line):
+            return output
+        if output is None:
+            output = []
+        output.append(PythonDirective(line_number, line_number, keyword, line))
+        return output
+
 
 def _parseKeywords(sourcepath, additional_parsers=[], require_script=True):
     """_parseKeywords
@@ -2022,6 +2161,7 @@
         IntegratedTestKeywordParser(
             "REDEFINE:", ParserKind.REDEFINE, initial_value=script
         ),
+        IntegratedTestKeywordParser("PYTHON:", ParserKind.PYTHON, initial_value=script),
     ]
     keyword_parsers = {p.keyword: p for p in builtin_parsers}
 
@@ -2047,9 +2187,11 @@
 
     # Verify the script contains a run line.
     if require_script and not any(
-        isinstance(directive, CommandDirective) for directive in script
+        isinstance(directive, CommandDirective)
+        or isinstance(directive, PythonDirective)
+        for directive in script
     ):
-        raise ValueError("Test has no 'RUN:' line")
+        raise ValueError("Test has no 'RUN:' line or 'PYTHON:' line")
 
     # Check for unterminated run or subst lines.
     #
@@ -2106,6 +2248,7 @@
     script = parsed["RUN:"] or []
     assert parsed["DEFINE:"] == script
     assert parsed["REDEFINE:"] == script
+    assert parsed["PYTHON:"] == script
     test.xfails += parsed["XFAIL:"] or []
     test.requires += parsed["REQUIRES:"] or []
     test.unsupported += parsed["UNSUPPORTED:"] or []
@@ -2142,56 +2285,190 @@
     return script
 
 
+# This is used in the lit test suite to check that symbols from lit's
+# implementation are not accidentally made available to 'PYTHON:' directives.
+def pythonDirectiveInaccessibleFunction():
+    assert False, "should be inaccessible from 'PYTHON:' directives"
+
+
+# script must be a list in which each item is either a (1) string that is a
+# shell command or (2) an ExpandableScriptDirective object, which might produce
+# shell commands.
+#
+# If substitutionApplier is not specified, lit substitutions in shell commands
+# are treated as plain text.  Otherwise substitutionApplier must be a function
+# that accepts a list of strings and returns the same list with lit
+# substitutions expanded.
+#
 # Always returns a lit.Test.Result.
-def _runShTest(test, litConfig, useExternalSh, script, tmpBase):
+def _runShTest(
+    test,
+    litConfig,
+    useExternalSh,
+    script,
+    tmpBase,
+    substitutionApplier=lambda x: x,
+):
+    # Make one attempt to run the test.
+    #
     # Always returns the tuple (out, err, exitCode, timeoutInfo).
     def runOnce(execdir):
-        # script is modified below (for litConfig.per_test_coverage, and for
-        # %dbg expansions).  runOnce can be called multiple times, but applying
-        # the modifications multiple times can corrupt script, so always modify
-        # a copy.
-        scriptCopy = script[:]
-        # Set unique LLVM_PROFILE_FILE for each run command
-        if litConfig.per_test_coverage:
-            # Extract the test case name from the test object, and remove the
-            # file extension.
-            test_case_name = test.path_in_suite[-1]
-            test_case_name = test_case_name.rsplit(".", 1)[0]
-            coverage_index = 0  # Counter for coverage file index
-            for i, ln in enumerate(scriptCopy):
-                match = re.fullmatch(kPdbgRegex, ln)
-                if match:
-                    dbg = match.group(1)
-                    command = match.group(2)
-                else:
-                    command = ln
-                profile = f"{test_case_name}{coverage_index}.profraw"
-                coverage_index += 1
-                command = f"export LLVM_PROFILE_FILE={profile}; {command}"
-                if match:
-                    command = buildPdbgCommand(dbg, command)
-                scriptCopy[i] = command
-
-        try:
+        # Counter for litConfig.per_test_coverage file names.
+        perTestCoverageIndex = 0
+
+        # scriptPart has the same constraints as script for _runShTest except
+        # scriptPart cannot contain python directives.
+        def runShScript(scriptPart, addStdout, addStderr):
+            scriptPart = substitutionApplier(scriptPart)
+            # If there were only substitution directives, there's nothing left now.
+            if not scriptPart:
+                return
+            # Set unique LLVM_PROFILE_FILE for each run command
+            if litConfig.per_test_coverage:
+                # Extract the test case name from the test object, and remove the
+                # file extension.
+                testCaseName = test.path_in_suite[-1]
+                testCaseName = testCaseName.rsplit(".", 1)[0]
+                for i, ln in enumerate(scriptPart):
+                    match = re.fullmatch(kPdbgRegex, ln)
+                    if match:
+                        dbg = match.group(1)
+                        command = match.group(2)
+                    else:
+                        command = ln
+                    nonlocal perTestCoverageIndex
+                    profile = f"{testCaseName}{perTestCoverageIndex}.profraw"
+                    perTestCoverageIndex += 1
+                    command = f"export LLVM_PROFILE_FILE={profile}; {command}"
+                    if match:
+                        command = buildPdbgCommand(dbg, command)
+                    scriptPart[i] = command
+            # Execute the remaining shell commands.
             if useExternalSh:
-                res = executeScript(test, litConfig, tmpBase, scriptCopy, execdir)
+                res = executeScript(test, litConfig, tmpBase, scriptPart, execdir)
             else:
                 res = executeScriptInternal(
-                    test, litConfig, tmpBase, scriptCopy, execdir
+                    test, litConfig, tmpBase, scriptPart, execdir
+                )
+            out, err, exitCode, timeoutInfo = res
+            addStdout(out)
+            addStderr(err)
+            if exitCode != 0:
+                raise ScriptFail(exitCode, timeoutInfo)
+
+        class PythonDirectiveLitAPI(object):
+            """
+            Class for the 'lit' object accessible in 'PYTHON:' directives and in
+            config.prologue.
+            """
+
+            # All members of PythonDirectiveLitAPI are exposed within 'PYTHON:'
+            # directives.  Symbols only required for lit iternals should be declared
+            # outside it.
+
+            # TODO: In the future, extend the 'lit' API to be able to *write*
+            # substitutions?
+
+            @staticmethod
+            def has(feature):
+                """
+                Check if config.available_features indicates a feature is enabled.
+                """
+                return feature in test.config.available_features
+
+            @staticmethod
+            def expand(text):
+                """
+                Expand substitutions in 'text' with their most recent values (from,
+                for example, the most recent 'DEFINE:' and 'REDEFINE:' directives),
+                and return the result.
+                """
+                return substitutionApplier([text])[0]
+
+            @staticmethod
+            def run(cmd):
+                """
+                Execute 'cmd' as if it appeared in a 'RUN:' directive.
+                """
+                msg = ""
+                # In the execution trace, report the call stack for this lit.run
+                # call.  However, skip this stack frame, and then climb the call
+                # stack, adding each frame, until we've reached this file again,
+                # and skip the rest.  That is, bothering the lit user with lit
+                # internals is just a distraction.
+                for caller in inspect.stack()[1:]:
+                    if caller.filename == __file__:
+                        break
+                    loc = f"{caller.filename}:{caller.lineno}"
+                    if not msg:
+                        msg += f"# lit.run called from {loc}"
+                    else:
+                        msg += "\n"
+                        msg += f"#         called from {loc}"
+                print(msg)
+                runShScript(
+                    # Empty "%dbg()" as the above suffices.
+                    [buildPdbgCommand("", cmd)],
+                    lambda s: print(s, end=""),
+                    lambda s: print(s, end="", file=sys.stderr),
                 )
-        except ScriptFatal as e:
-            out = f"# " + "\n# ".join(str(e).splitlines()) + "\n"
-            return out, "", 1, None, Test.UNRESOLVED
 
-        out, err, exitCode, timeoutInfo = res
-        if exitCode == 0:
-            status = Test.PASS
-        else:
-            if timeoutInfo is None:
-                status = Test.FAIL
-            else:
-                status = Test.TIMEOUT
-        return out, err, exitCode, timeoutInfo, status
+        out = err = ""
+
+        def addStdout(text):
+            nonlocal out
+            out += text
+
+        def addStderr(text):
+            nonlocal err
+            err += text
+
+        python_dict = dict()
+        python_dict["lit"] = PythonDirectiveLitAPI
+        try:
+            # Execute the python prologue, if any.
+            prologue = test.config.prologue
+            if prologue:
+                with open(prologue) as f:
+                    what = f"config.prologue='{prologue}'"
+                    PythonDirective.executePython(
+                        what, prologue, f.read(), python_dict, addStdout
+                    )
+            # Execute the script.
+            scriptRemaining = script
+            while scriptRemaining:
+                # Execute all leading python directives.
+                for pythonDirIndex, pythonDir in enumerate(scriptRemaining):
+                    if not isinstance(pythonDir, PythonDirective):
+                        break
+                    pythonDir.execute(test.getSourcePath(), python_dict, addStdout)
+                else:
+                    scriptRemaining = []
+                    break
+                scriptRemaining = scriptRemaining[pythonDirIndex:]
+                # Extract all directives before the next python directive into
+                # a script fragment.
+                pythonDirIndex = next(
+                    (
+                        i
+                        for i, d in enumerate(scriptRemaining)
+                        if isinstance(d, PythonDirective)
+                    ),
+                    len(scriptRemaining),
+                )
+                scriptPart = scriptRemaining[:pythonDirIndex]
+                scriptRemaining = scriptRemaining[pythonDirIndex:]
+                # In that script fragment, execute substitution directives and
+                # expand substitutions in shell commands, and then execute the
+                # resulting shell script.
+                runShScript(scriptPart, addStdout, addStderr)
+                # Repeat until nothing's left in the original script.
+        except ScriptFatal as e:
+            out += f"# " + "\n# ".join(str(e).splitlines()) + "\n"
+            return out, err, 1, None, Test.UNRESOLVED
+        except ScriptFail as e:
+            return out, err, e.exitCode, e.timeoutInfo, e.status
+        return out, err, 0, None, Test.PASS
 
     # Create the output directory if it does not already exist.
     lit.util.mkdir_p(os.path.dirname(tmpBase))
@@ -2249,11 +2526,20 @@
         test, tmpDir, tmpBase, normalize_slashes=useExternalSh
     )
     conditions = {feature: True for feature in test.config.available_features}
-    script = applySubstitutions(
+
+    def substitutionApplier(lines):
+        return applySubstitutions(
+            lines,
+            substitutions,
+            conditions,
+            recursion_limit=test.config.recursiveExpansionLimit,
+        )
+
+    return _runShTest(
+        test,
+        litConfig,
+        useExternalSh,
         script,
-        substitutions,
-        conditions,
-        recursion_limit=test.config.recursiveExpansionLimit,
+        tmpBase,
+        substitutionApplier=substitutionApplier,
     )
-
-    return _runShTest(test, litConfig, useExternalSh, script, tmpBase)
Index: llvm/utils/lit/lit/TestingConfig.py
===================================================================
--- llvm/utils/lit/lit/TestingConfig.py
+++ llvm/utils/lit/lit/TestingConfig.py
@@ -192,6 +192,7 @@
         self.limit_to_features = set(limit_to_features)
         self.parallelism_group = parallelism_group
         self._recursiveExpansionLimit = None
+        self.prologue = None
 
     @property
     def recursiveExpansionLimit(self):
Index: llvm/utils/lit/tests/Inputs/per-test-coverage-by-lit-cfg/per-test-coverage-by-lit-cfg.py
===================================================================
--- llvm/utils/lit/tests/Inputs/per-test-coverage-by-lit-cfg/per-test-coverage-by-lit-cfg.py
+++ llvm/utils/lit/tests/Inputs/per-test-coverage-by-lit-cfg/per-test-coverage-by-lit-cfg.py
@@ -1,6 +1,7 @@
 # Check that the environment variable is set correctly
 # RUN: %{python} %s | FileCheck -DINDEX=1 %s
-# RUN: %{python} %s | FileCheck -DINDEX=2 %s
+# PYTHON: lit.run("%{python} %s | FileCheck -DINDEX=2 %s")
+# RUN: %{python} %s | FileCheck -DINDEX=3 %s
 
 # Python script to read the environment variable
 # and print its value
Index: llvm/utils/lit/tests/Inputs/per-test-coverage/per-test-coverage.py
===================================================================
--- llvm/utils/lit/tests/Inputs/per-test-coverage/per-test-coverage.py
+++ llvm/utils/lit/tests/Inputs/per-test-coverage/per-test-coverage.py
@@ -1,6 +1,7 @@
 # Check that the environment variable is set correctly
 # RUN: %{python} %s | FileCheck -DINDEX=1 %s
-# RUN: %{python} %s | FileCheck -DINDEX=2 %s
+# PYTHON: lit.run("%{python} %s | FileCheck -DINDEX=2 %s")
+# RUN: %{python} %s | FileCheck -DINDEX=3 %s
 
 # Python script to read the environment variable
 # and print its value
Index: llvm/utils/lit/tests/Inputs/shtest-define/errors/no-run.txt
===================================================================
--- llvm/utils/lit/tests/Inputs/shtest-define/errors/no-run.txt
+++ llvm/utils/lit/tests/Inputs/shtest-define/errors/no-run.txt
@@ -3,6 +3,6 @@
 # DEFINE: %{local:echo}=foo
 # REDEFINE: %{global:echo}=bar
 
-# CHECK: Test has no '{{RUN}}:' line
+# CHECK: Test has no '{{RUN}}:' line or '{{PYTHON}}:' line
 
 # CHECK: Unresolved: 1
Index: llvm/utils/lit/tests/Inputs/shtest-python/errors/README.md
===================================================================
--- /dev/null
+++ llvm/utils/lit/tests/Inputs/shtest-python/errors/README.md
@@ -0,0 +1,37 @@
+One goal of the tests in this directory is to check that diagnostics from the
+python interpreter are faithfully copied to the lit user.  However, verifying
+the exact text of those diagnostics makes the tests brittle across python
+versions.
+
+Please follow the guide below when adding new tests, and update it if you find
+additional issues.  If you think of a better strategy for checking lit's
+handling of python diagnostics, please propose it.
+
+So far in the tests we have written, the following diagnostic text seems to be
+stable across python versions, and it seems worthwhile to verify that it is
+copied faithfully to the lit user:
+
+  - The traceback.
+  - When the erroneous python statement contains only one line, it is the last
+    line that is quoted from the PYTHON block.
+  - The presence of a caret (`^`) after the last line quoted from the PYTHON
+    block.
+  - The exception, if raised directly by our test code.
+
+The following have proven unstable:
+
+  - The last line that is quoted from the PYTHON block when the erroneous
+    python statement contains multiple lines.
+  - The number of context lines quoted from the PYTHON block.  Of course, if
+    the PYTHON block has only one line, that number is stably zero.
+  - The indentation of the last line quoted from the PYTHON block.  In
+    particular, we have found that, independently of lit, at least python 3.8.0
+    sometimes botches the indentation relative to other quoted lines and the
+    caret, making the diagnostic confusing.
+  - The indentation of the caret, including where it points in the last quoted
+    line.
+  - The exception, including its kind (e.g., `SyntaxError: invalid syntax`), if
+    not raised directly by our test code.
+  - In a compilation error, whether a line quoted from a PYTHON block includes
+    only the python code or the full source line including the `# PYTHON: `
+    prefix.
Index: llvm/utils/lit/tests/Inputs/shtest-python/errors/incomplete-python.txt
===================================================================
--- /dev/null
+++ llvm/utils/lit/tests/Inputs/shtest-python/errors/incomplete-python.txt
@@ -0,0 +1,17 @@
+# Incomplete python code is a special error case that some tools for compiling
+# python (e.g., code.compile_command) leave to the caller.  Whether or not lit
+# uses such tools now, make sure lit doesn't start to miss this case.
+
+# PYTHON: if True: #LN
+
+#       CHECK:Command Output (stdout):
+#  CHECK-NEXT:--
+#  CHECK-NEXT:# error compiling '{{PYTHON}}:' directive at line [[#LN: @LINE - 4]]:
+#  CHECK-NEXT:#   File "[[TEST_FILE]]", line [[#LN]]
+#  CHECK-NEXT:# {{ *(# PYTHON[:] )?}}if True: #LN
+#  CHECK-NEXT:# {{ *}}^
+#  CHECK-NEXT:# {{.*Error.*}}
+# CHECK-EMPTY:
+#  CHECK-NEXT:--
+
+# CHECK:  Unresolved: 1
Index: llvm/utils/lit/tests/Inputs/shtest-python/errors/inconsistent-indent-2-lines.txt
===================================================================
--- /dev/null
+++ llvm/utils/lit/tests/Inputs/shtest-python/errors/inconsistent-indent-2-lines.txt
@@ -0,0 +1,7 @@
+# PYTHON:  3 + \
+# PYTHON: 4
+
+#      CHECK:'{{PYTHON}}:' directive continuation line has indentation that is inconsistent with previous lines
+# CHECK-NEXT:in {{PYTHON}}: directive on test line 2
+
+# CHECK:  Unresolved: 1
Index: llvm/utils/lit/tests/Inputs/shtest-python/errors/inconsistent-indent-3-lines.txt
===================================================================
--- /dev/null
+++ llvm/utils/lit/tests/Inputs/shtest-python/errors/inconsistent-indent-3-lines.txt
@@ -0,0 +1,8 @@
+# PYTHON: 3 + \
+# PYTHON: 4 *  \
+# PYTHON:5
+
+#      CHECK:'{{PYTHON}}:' directive continuation line has indentation that is inconsistent with previous lines
+# CHECK-NEXT:in {{PYTHON}}: directive on test line 3
+
+# CHECK:  Unresolved: 1
Index: llvm/utils/lit/tests/Inputs/shtest-python/errors/internal-api-inaccessible.txt
===================================================================
--- /dev/null
+++ llvm/utils/lit/tests/Inputs/shtest-python/errors/internal-api-inaccessible.txt
@@ -0,0 +1,19 @@
+# pythonDirectiveInaccessibleFunction is defined within lit's implementation
+# (specifically for the purposes of the test).  It and other lit implementation
+# symbols should not be accessible to PYTHON directives.
+
+# PYTHON: pythonDirectiveInaccessibleFunction() # LN
+
+#       CHECK:Command Output (stdout):
+#  CHECK-NEXT:--
+#  CHECK-NEXT:# executed '{{PYTHON}}:' directive at line [[#LN: @LINE - 4]]
+#  CHECK-NEXT:# .---stderr from '{{PYTHON}}:' directive at line [[#LN]]
+#  CHECK-NEXT:# | Traceback{{.*}}:
+#  CHECK-NEXT:# |   File "[[TEST_FILE]]", line [[#LN]], in <module>
+#  CHECK-NEXT:# | {{ *}}# {{PYTHON}}: pythonDirectiveInaccessibleFunction() # LN
+#  CHECK-NEXT:# | {{.*Error.*pythonDirectiveInaccessibleFunction.*}}
+#  CHECK-NEXT:# `-----------------------------
+# CHECK-EMPTY:
+#  CHECK-NEXT:--
+
+# CHECK:  Failed: 1
Index: llvm/utils/lit/tests/Inputs/shtest-python/errors/lit.prologue.py
===================================================================
--- /dev/null
+++ llvm/utils/lit/tests/Inputs/shtest-python/errors/lit.prologue.py
@@ -0,0 +1,6 @@
+import inspect, sys
+
+print("config.prologue writes to stdout")
+print("config.prologue writes to stderr", file=sys.stderr)
+line = str(inspect.stack()[0].lineno + 1)
+raise Exception(f"exception in config.prologue at line {line}")
Index: llvm/utils/lit/tests/Inputs/shtest-python/errors/python-syntax-error-1-line.txt
===================================================================
--- /dev/null
+++ llvm/utils/lit/tests/Inputs/shtest-python/errors/python-syntax-error-1-line.txt
@@ -0,0 +1,13 @@
+# PYTHON:+
+
+#       CHECK:Command Output (stdout):
+#  CHECK-NEXT:--
+#  CHECK-NEXT:# error compiling '{{PYTHON}}:' directive at line 1:
+#  CHECK-NEXT:#   File "[[TEST_FILE]]", line 1
+#  CHECK-NEXT:# {{ *(# PYTHON[:])?}}+
+#  CHECK-NEXT:# {{ *}}^
+#  CHECK-NEXT:# {{.*Error.*}}
+# CHECK-EMPTY:
+#  CHECK-NEXT:--
+
+# CHECK:  Unresolved: 1
Index: llvm/utils/lit/tests/Inputs/shtest-python/errors/python-syntax-error-2-lines.txt
===================================================================
--- /dev/null
+++ llvm/utils/lit/tests/Inputs/shtest-python/errors/python-syntax-error-2-lines.txt
@@ -0,0 +1,14 @@
+# PYTHON:3 +   \
+# PYTHON:  * 4
+
+#       CHECK:Command Output (stdout):
+#  CHECK-NEXT:--
+#  CHECK-NEXT:# error compiling '{{PYTHON}}:' directive from line 1 to 2:
+#  CHECK-NEXT:#   File "[[TEST_FILE]]", line 2
+#       CHECK:# {{ *(# PYTHON[:]  )?}}* 4
+#  CHECK-NEXT:# {{ *}}^
+#  CHECK-NEXT:# {{.*Error.*}}
+# CHECK-EMPTY:
+#  CHECK-NEXT:--
+
+# CHECK:  Unresolved: 1
Index: llvm/utils/lit/tests/Inputs/shtest-python/errors/python-syntax-error-3-lines.txt
===================================================================
--- /dev/null
+++ llvm/utils/lit/tests/Inputs/shtest-python/errors/python-syntax-error-3-lines.txt
@@ -0,0 +1,18 @@
+# PYTHON:3 +   \
+# PYTHON:  4 * \
+# PYTHON:  / 5
+
+# Some versions of python (e.g., 3.8.0) botch the indentation on line
+# continuations, so we accept any indentation there for now.
+
+#       CHECK:Command Output (stdout):
+#  CHECK-NEXT:--
+#  CHECK-NEXT:# error compiling '{{PYTHON}}:' directive from line 1 to 3:
+#  CHECK-NEXT:#   File "[[TEST_FILE]]", line 3
+#       CHECK:# {{ *(# PYTHON[:]  )?}}/ 5
+#  CHECK-NEXT:# {{ *}}^
+#  CHECK-NEXT:# {{.*Error.*}}
+# CHECK-EMPTY:
+#  CHECK-NEXT:--
+
+# CHECK:  Unresolved: 1
Index: llvm/utils/lit/tests/Inputs/shtest-python/errors/python-syntax-error-leading-indent.txt
===================================================================
--- /dev/null
+++ llvm/utils/lit/tests/Inputs/shtest-python/errors/python-syntax-error-leading-indent.txt
@@ -0,0 +1,18 @@
+# PYTHON:  3 +   \
+# PYTHON:    4 * \
+# PYTHON:    / 5
+
+# Some versions of python (e.g., 3.8.0) botch the indentation on line
+# continuations, so we accept any indentation there for now.
+
+#       CHECK:Command Output (stdout):
+#  CHECK-NEXT:--
+#  CHECK-NEXT:# error compiling '{{PYTHON}}:' directive from line 1 to 3:
+#  CHECK-NEXT:#   File "[[TEST_FILE]]", line 3
+#       CHECK:# {{ *(# PYTHON[:]    )?}}/ 5
+#  CHECK-NEXT:# {{ *}}^
+#  CHECK-NEXT:# {{.*Error.*}}
+# CHECK-EMPTY:
+#  CHECK-NEXT:--
+
+# CHECK:  Unresolved: 1
Index: llvm/utils/lit/tests/Inputs/shtest-python/errors/trace-compile-error.txt
===================================================================
--- /dev/null
+++ llvm/utils/lit/tests/Inputs/shtest-python/errors/trace-compile-error.txt
@@ -0,0 +1,38 @@
+# Check that the execution trace is correct when a PYTHON directive fails to
+# compile.  Specifically, make sure the exception is written at the end of the
+# execution trace and doesn't prevent stdout and stderr from prior RUN
+# directives and python code from appearing.
+
+# PYTHON: import sys                                        #LN
+# PYTHON: print("PYTHON writes to stdout")                  #LN + 1
+# PYTHON: print("PYTHON writes to stderr", file=sys.stderr) #LN + 2
+#    RUN: %{python} %S/../write-to-stdout-and-stderr.py
+# PYTHON: +
+
+#       CHECK:Command Output (stdout):
+#  CHECK-NEXT:--
+#  CHECK-NEXT:# executed '{{PYTHON}}:' directive from line [[#LN : @LINE - 8]] to [[# @LINE - 6]]
+#  CHECK-NEXT:# .---stdout from '{{PYTHON}}:' directive from line [[#LN]] to [[#LN + 2]]
+#  CHECK-NEXT:# | PYTHON writes to stdout
+#  CHECK-NEXT:# `-----------------------------
+#  CHECK-NEXT:# .---stderr from '{{PYTHON}}:' directive from line [[#LN]] to [[#LN + 2]]
+#  CHECK-NEXT:# | PYTHON writes to stderr
+#  CHECK-NEXT:# `-----------------------------
+#  CHECK-NEXT:# {{RUN}}: at line [[#LN + 3]]
+#  CHECK-NEXT:{{.*/python.*}} {{.*}}/write-to-stdout-and-stderr.py
+#  CHECK-NEXT:# executed command: {{.*/python.*}} {{.*}}/write-to-stdout-and-stderr.py
+#  CHECK-NEXT:# .---command stdout------------
+#  CHECK-NEXT:# | a line on stdout
+#  CHECK-NEXT:# `-----------------------------
+#  CHECK-NEXT:# .---command stderr------------
+#  CHECK-NEXT:# | a line on stderr
+#  CHECK-NEXT:# `-----------------------------
+#  CHECK-NEXT:# error compiling '{{PYTHON}}:' directive at line [[#LN + 4]]:
+#  CHECK-NEXT:#   File "[[TEST_FILE]]", line [[#LN + 4]]
+#  CHECK-NEXT:# {{ *(# PYTHON[:] )?}}+
+#  CHECK-NEXT:# {{ *}}^
+#  CHECK-NEXT:# {{.*Error.*}}
+# CHECK-EMPTY:
+#  CHECK-NEXT:--
+
+# CHECK:  Unresolved: 1
Index: llvm/utils/lit/tests/Inputs/shtest-python/errors/trace-exception.txt
===================================================================
--- /dev/null
+++ llvm/utils/lit/tests/Inputs/shtest-python/errors/trace-exception.txt
@@ -0,0 +1,37 @@
+# Check that the execution trace is correct when a PYTHON directive raises an
+# exception.  Specifically, make sure the exception is written at the end of the
+# associated stderr and doesn't prevent prior stdout and stderr from appearing.
+# Also, make sure non-directive lines mixed in with the PYTHON directives
+# doesn't throw off the line numbers in the python traceback.
+
+# PYTHON: import sys                                        #LN_START
+
+# Non-directive lines might be used for formatting or comments like this one.
+# PYTHON: print("PYTHON writes to stdout")
+# PYTHON: print("PYTHON writes to stderr", file=sys.stderr)
+#
+#
+# PYTHON: def fn():
+# PYTHON:     raise Exception("fail")
+#
+# PYTHON: fn()                                              #LN_END
+
+#       CHECK:Command Output (stdout):
+#  CHECK-NEXT:--
+#  CHECK-NEXT:# executed '{{PYTHON}}:' directive from line [[#LN_START: @LINE - 14]] to [[#LN_END: @LINE - 4]]
+#  CHECK-NEXT:# .---stdout from '{{PYTHON}}:' directive from line [[#LN_START]] to [[#LN_END]]
+#  CHECK-NEXT:# | PYTHON writes to stdout
+#  CHECK-NEXT:# `-----------------------------
+#  CHECK-NEXT:# .---stderr from '{{PYTHON}}:' directive from line [[#LN_START]] to [[#LN_END]]
+#  CHECK-NEXT:# | PYTHON writes to stderr
+#  CHECK-NEXT:# | Traceback{{.*}}:
+#  CHECK-NEXT:# |   File "[[TEST_FILE]]", line [[#LN_END]], in <module>
+#  CHECK-NEXT:# | {{ *}}# {{PYTHON}}: fn() {{ *}} #LN_END
+#  CHECK-NEXT:# |   File "[[TEST_FILE]]", line [[#LN_END - 2]], in fn
+#  CHECK-NEXT:# | {{ *}}# {{PYTHON}}:     raise Exception("fail")
+#  CHECK-NEXT:# | Exception: fail
+#  CHECK-NEXT:# `-----------------------------
+# CHECK-EMPTY:
+#  CHECK-NEXT:--
+
+# CHECK:  Failed: 1
Index: llvm/utils/lit/tests/Inputs/shtest-python/errors/trace-prologue-exception.txt
===================================================================
--- /dev/null
+++ llvm/utils/lit/tests/Inputs/shtest-python/errors/trace-prologue-exception.txt
@@ -0,0 +1,26 @@
+# Check that the execution trace is correct when config.prologue raises an
+# exception.  Specifically, make sure the exception is written at the end of the
+# associated stderr and doesn't prevent prior stdout and stderr from
+# config.prologue from appearing.  However, no PYTHON directive should execute
+# or produce any output.
+
+# PYTHON: print("PYTHON writes to stdout")                  #LN
+# PYTHON: print("PYTHON writes to stderr", file=sys.stderr) #LN + 2
+#
+#      CHECK:Command Output (stdout):
+# CHECK-NEXT:--
+#  CHECK-NEXT:# executed config.prologue='[[PROLOGUE:.*/lit.prologue.py]]'
+#  CHECK-NEXT:# .---stdout from config.prologue='[[PROLOGUE]]'
+#  CHECK-NEXT:# | config.prologue writes to stdout
+#  CHECK-NEXT:# `-----------------------------
+#  CHECK-NEXT:# .---stderr from config.prologue='[[PROLOGUE]]'
+#  CHECK-NEXT:# | config.prologue writes to stderr
+#  CHECK-NEXT:# | Traceback{{.*}}:
+#  CHECK-NEXT:# |   File "[[PROLOGUE]]", line [[#LOC_LN:]], in <module>
+#  CHECK-NEXT:# |     raise Exception({{.*}})
+#  CHECK-NEXT:# | Exception: exception in config.prologue at line [[#LOC_LN]]
+#  CHECK-NEXT:# `-----------------------------
+# CHECK-EMPTY:
+#  CHECK-NEXT:--
+
+# CHECK:  Failed: 1
Index: llvm/utils/lit/tests/Inputs/shtest-python/errors/trace-run-compile-error.txt
===================================================================
--- /dev/null
+++ llvm/utils/lit/tests/Inputs/shtest-python/errors/trace-run-compile-error.txt
@@ -0,0 +1,39 @@
+# Check that the execution trace is correct when a RUN directive fails to
+# compile.  Specifically, make sure the error is written at the end of the
+# execution trace and doesn't prevent stdout and stderr from prior RUN
+# directives and python code from appearing.
+#
+# Before PYTHON directives were introduced to lit, this scenario was an
+# impossible case.  There were no PYTHON directives to break RUN directives into
+# multiple scripts, so all RUN directives were compiled before any were
+# executed, so there could not have been prior stdout or stderr.
+
+#    RUN: %{python} %S/../write-to-stdout-and-stderr.py
+# PYTHON: import sys                                        #LN + 1
+# PYTHON: print("PYTHON writes to stdout")                  #LN + 2
+# PYTHON: print("PYTHON writes to stderr", file=sys.stderr) #LN + 3
+#    RUN: &&
+
+#       CHECK:Command Output (stdout):
+#  CHECK-NEXT:--
+#  CHECK-NEXT:# {{RUN}}: at line [[#LN: @LINE - 8]]
+#  CHECK-NEXT:{{.*/python.*}} {{.*}}/write-to-stdout-and-stderr.py
+#  CHECK-NEXT:# executed command: {{.*/python.*}} {{.*}}/write-to-stdout-and-stderr.py
+#  CHECK-NEXT:# .---command stdout------------
+#  CHECK-NEXT:# | a line on stdout
+#  CHECK-NEXT:# `-----------------------------
+#  CHECK-NEXT:# .---command stderr------------
+#  CHECK-NEXT:# | a line on stderr
+#  CHECK-NEXT:# `-----------------------------
+#  CHECK-NEXT:# executed '{{PYTHON}}:' directive from line [[#LN + 1]] to [[#LN + 3]]
+#  CHECK-NEXT:# .---stdout from '{{PYTHON}}:' directive from line [[#LN + 1]] to [[#LN + 3]]
+#  CHECK-NEXT:# | PYTHON writes to stdout
+#  CHECK-NEXT:# `-----------------------------
+#  CHECK-NEXT:# .---stderr from '{{PYTHON}}:' directive from line [[#LN + 1]] to [[#LN + 3]]
+#  CHECK-NEXT:# | PYTHON writes to stderr
+#  CHECK-NEXT:# `-----------------------------
+#  CHECK-NEXT:# shell parser error on RUN: at line [[#LN + 4]]: &&
+# CHECK-EMPTY:
+#  CHECK-NEXT:--
+
+# CHECK:  Unresolved: 1
Index: llvm/utils/lit/tests/Inputs/shtest-python/exampleModule.py
===================================================================
--- /dev/null
+++ llvm/utils/lit/tests/Inputs/shtest-python/exampleModule.py
@@ -0,0 +1,15 @@
+# The importer should assign this to the 'lit' object, or uses of 'lit' below
+# will fail.
+lit = None
+
+import inspect
+
+
+def helloWorldFromLitRun():
+    loc = f"{inspect.stack()[0].filename}:{inspect.stack()[0].lineno + 1}"
+    lit.run(f"echo hello world from lit.run at {loc}")
+
+
+def goodbyeWorldFromLitRun():
+    loc = f"{inspect.stack()[0].filename}:{inspect.stack()[0].lineno + 1}"
+    lit.run(f"echo goodbye world from lit.run at {loc}")
Index: llvm/utils/lit/tests/Inputs/shtest-python/file.txt
===================================================================
--- /dev/null
+++ llvm/utils/lit/tests/Inputs/shtest-python/file.txt
@@ -0,0 +1,5 @@
+# Check that __file__ is defined in PYTHON lines.
+
+# PYTHON: print(f"test file <{__file__}>")
+
+# CHECK:# | test file <[[TEST_FILE]]>
Index: llvm/utils/lit/tests/Inputs/shtest-python/has.txt
===================================================================
--- /dev/null
+++ llvm/utils/lit/tests/Inputs/shtest-python/has.txt
@@ -0,0 +1,4 @@
+# PYTHON: assert lit.has('enabledFeature')
+# PYTHON: assert not lit.has('disabledFeature')
+
+# CHECK:  Passed: 1
Index: llvm/utils/lit/tests/Inputs/shtest-python/import.txt
===================================================================
--- /dev/null
+++ llvm/utils/lit/tests/Inputs/shtest-python/import.txt
@@ -0,0 +1,48 @@
+# For safety, lit doesn't automatically add %S to sys.path.
+#
+# PYTHON: try:
+# PYTHON:     import exampleModule
+# PYTHON: except ModuleNotFoundError:
+# PYTHON:     pass
+# PYTHON: else:
+# PYTHON:     raise Exception("error: sys.path already contains %S")
+
+# Adding %S to sys.path is easy enough from a PYTHON directive.
+#
+# PYTHON: import sys
+# PYTHON: sys.path.append(lit.expand('%S'))
+# PYTHON: import exampleModule
+
+# lit doesn't automatically expose the 'lit' object to an imported module, which
+# might have nothing to do with lit.
+#
+# PYTHON: try:
+# PYTHON:     exampleModule.helloWorldFromLitRun()
+# PYTHON: except AttributeError:
+# PYTHON:     pass
+# PYTHON: else:
+# PYTHON:     raise Exception("error: imported module already sees 'lit'")
+
+# Every function in an imported module could be written to accept the 'lit'
+# object as a parameter, but it seems easier to just expose it once via a PYTHON
+# directive as follows.
+#
+# PYTHON: exampleModule.lit = lit
+# PYTHON: exampleModule.helloWorldFromLitRun()
+# PYTHON: exampleModule.goodbyeWorldFromLitRun()
+#
+# CHECK:# | # | hello world from lit.run at {{.*}}
+# CHECK:# | # | goodbye world from lit.run at {{.*}}
+
+# In summary, here's the boilerplate to set up use of exampleModule in a single
+# test:
+#
+#   import sys
+#   sys.path.append(lit.expand('%S'))
+#   import exampleModule
+#   exampleModule.lit = lit
+#
+# If you have to do this in every test in a test suite, it's easier to use
+# config.prologue.  See the test ./prologue.txt for an example.
+
+# CHECK:  Passed: 1
Index: llvm/utils/lit/tests/Inputs/shtest-python/lit.cfg
===================================================================
--- /dev/null
+++ llvm/utils/lit/tests/Inputs/shtest-python/lit.cfg
@@ -0,0 +1,17 @@
+import lit.formats
+
+config.name = "shtest-python"
+config.suffixes = [".txt"]
+# Use lit's internal shell to avoid shell portability issues within RUN lines
+# (e.g., for 'echo' commands in Windows).  Those issues should be orthogonal to
+# the python directive behavior we are trying to test.
+config.test_format = lit.formats.ShTest(execute_external=False)
+config.test_source_root = None
+config.test_exec_root = None
+config.substitutions.append(("%{python}", f'"{sys.executable}"'))
+
+config.available_features.add("enabledFeature")
+
+prologue = lit_config.params.get("addPrologue", None)
+if prologue:
+    config.prologue = os.path.join(os.path.dirname(__file__), prologue)
Index: llvm/utils/lit/tests/Inputs/shtest-python/lit.prologue.py
===================================================================
--- /dev/null
+++ llvm/utils/lit/tests/Inputs/shtest-python/lit.prologue.py
@@ -0,0 +1,39 @@
+# If lit.cfg sets config.prologue to this file, this file executes immediately
+# before the first PYTHON directive almost if it appears in a PYTHON directive.
+# One notable difference is __file__ is set to the absolute location of this
+# file to help with imports.
+
+import os, sys
+
+sys.path.append(os.path.dirname(__file__))
+
+import exampleModule
+
+exampleModule.lit = lit  # so the module can access the lit object
+
+# The following helps us test the execution trace of config.prologue.
+
+import inspect
+
+print("config.prologue writes to stdout")
+print("config.prologue writes to stderr", file=sys.stderr)
+lit.run("true")  # writes nothing to stdout or stderr
+loc = str(inspect.stack()[0].lineno + 1)
+lit.run("echo config.prologue lit.run writes to stdout at line " + loc)
+loc = str(inspect.stack()[0].lineno + 1)
+lit.run("%{python} %S/write-to-stderr.py && echo at line " + loc)
+loc = str(inspect.stack()[0].lineno + 1)
+lit.run("%{python} %S/write-to-stdout-and-stderr.py && echo at line " + loc)
+
+
+# fmt: off
+def localFnA():
+    loc = str(inspect.stack()[0].lineno + 1)
+    lit.run("echo config.prologue lit.run from func at line " + loc)
+def localFnB():
+    localFnA()  # must be at above lit.run line number + 2
+localFnB()  # must be at above lit.run line number + 3
+# fmt: on
+exampleModule.helloWorldFromLitRun()
+
+del inspect, os, sys  # optional: just avoids unnecessary pollution across tests
Index: llvm/utils/lit/tests/Inputs/shtest-python/no-shell-commands.txt
===================================================================
--- /dev/null
+++ llvm/utils/lit/tests/Inputs/shtest-python/no-shell-commands.txt
@@ -0,0 +1,12 @@
+# Check the case where the entire test file executes python but no shell
+# commands.  This should be fine.
+
+# PYTHON: pass
+
+#       CHECK:Command Output (stdout):
+#  CHECK-NEXT:--
+#  CHECK-NEXT:# executed '{{PYTHON}}:' directive at line [[# @LINE - 4]]
+# CHECK-EMPTY:
+#  CHECK-NEXT:--
+
+# CHECK:  Passed: 1
Index: llvm/utils/lit/tests/Inputs/shtest-python/prologue.txt
===================================================================
--- /dev/null
+++ llvm/utils/lit/tests/Inputs/shtest-python/prologue.txt
@@ -0,0 +1,32 @@
+# To avoid the boilerplate discussed in the test ./import.txt to set up use of a
+# module, this test depends on lit.cfg's config.prologue.
+#
+# That config.prologue has writes to stdout/stderr so we can check when it's
+# executed:
+# - It should be executed before any RUN or PYTHON directive.  This test checks
+#   the case when the first directive is a RUN directive.  ./trace.txt checks
+#   the case when the first directive is a PYTHON directive.
+# - It should be executed only once even if there are multiple PYTHON blocks
+#   and regardless of the number of lines in those blocks.
+
+#    RUN: echo after config.prologue executes
+# PYTHON: exampleModule.helloWorldFromLitRun()
+#    RUN: echo start new python block
+# PYTHON: exampleModule.helloWorldFromLitRun()
+# PYTHON: exampleModule.goodbyeWorldFromLitRun()
+
+# CHECK-NOT:config.prologue writes
+#     CHECK:# | config.prologue writes to stdout
+# CHECK-NOT:config.prologue writes
+#     CHECK:# | config.prologue writes to stderr
+# CHECK-NOT:config.prologue writes
+#     CHECK:# | after config.prologue executes
+# CHECK-NOT:config.prologue writes
+#     CHECK:# | # | hello world from lit.run at {{.*}}
+# CHECK-NOT:config.prologue writes
+#     CHECK:# | # | hello world from lit.run at {{.*}}
+# CHECK-NOT:config.prologue writes
+#     CHECK:# | # | goodbye world from lit.run at {{.*}}
+# CHECK-NOT:config.prologue writes
+
+# CHECK:  Passed: 1
Index: llvm/utils/lit/tests/Inputs/shtest-python/shell-affects-python.txt
===================================================================
--- /dev/null
+++ llvm/utils/lit/tests/Inputs/shtest-python/shell-affects-python.txt
@@ -0,0 +1,23 @@
+# Check that the effect of a RUN directive or a lit.run call in a PYTHON
+# directive can be seen by later python code in PYTHON directives.  That is, the
+# role of PYTHON directives is more than just calling lit.run to insert
+# additional shell commands while building the RUN directives' shell script for
+# later execution.  Instead, PYTHON directives must execute as part of that
+# script.
+
+#    RUN: mkdir -p %t
+# PYTHON: lit.run('touch %t/python-different-block')
+#    RUN: touch %t/run
+
+# PYTHON: import os
+# PYTHON: lit.run('touch %t/python-same-block')
+# PYTHON: for file in sorted(os.listdir(lit.expand('%t'))):
+# PYTHON:   print(file)
+
+# CHECK-COUNT-3:{{.*}} executed command: touch {{.*}}
+#  CHECK-NEXT:# | python-different-block
+#  CHECK-NEXT:# | python-same-block
+#  CHECK-NEXT:# | run
+#  CHECK-NEXT:# `--{{-*}}
+
+# CHECK:  Passed: 1
Index: llvm/utils/lit/tests/Inputs/shtest-python/substs-affects-python.txt
===================================================================
--- /dev/null
+++ llvm/utils/lit/tests/Inputs/shtest-python/substs-affects-python.txt
@@ -0,0 +1,25 @@
+# Check that the effect of DEFINE/REDEFINE directives can be seen by later
+# PYTHON directives' lit.expand or lit.run.
+
+# PYTHON: def run():
+# PYTHON:     lit.run('echo %{greeting} from lit.run')
+# PYTHON:     greeting = lit.expand('%{greeting}')
+# PYTHON:     print(f"{greeting} from lit.expand")
+
+# DEFINE: %{greeting} = hello
+# PYTHON: run()
+
+# REDEFINE: %{greeting} = goodbye
+# PYTHON: run()
+
+# CHECK-NOT:{{^(hello|goodbye)}}
+#     CHECK:# | # | hello from lit.run
+# CHECK-NOT:{{^(hello|goodbye)}}
+#     CHECK:# | hello from lit.expand
+# CHECK-NOT:{{^(hello|goodbye)}}
+#     CHECK:# | # | goodbye from lit.run
+# CHECK-NOT:{{^(hello|goodbye)}}
+#     CHECK:# | goodbye from lit.expand
+# CHECK-NOT:{{^(hello|goodbye)}}
+
+# CHECK:  Passed: 1
Index: llvm/utils/lit/tests/Inputs/shtest-python/trace.txt
===================================================================
--- /dev/null
+++ llvm/utils/lit/tests/Inputs/shtest-python/trace.txt
@@ -0,0 +1,242 @@
+# The execution trace is captured as stdout under the header "Command Output
+# (stdout)".  This test uses lit's internal shell, which doesn't use the header
+# "Command Output (stderr)".  (When using an external shell, stderr from shell
+# commands is handled differently.)
+#
+#      CHECK:Command Output (stdout):
+# CHECK-NEXT:--
+
+# Execution traces behave as follows.  First, a comment is added to indicate the
+# start of execution of each of the following, and that comment is followed by a
+# header for all captured stdout and a separate header for all captured stderr,
+# if any:
+#
+# - config.prologue.
+# - PYTHON directive block (includes any consecutive PYTHON directives).
+# - lit.run call.  Its trace appears in the stdout of the associated
+#   config.prologue or PYTHON directive block.  Its starting comment includes
+#   locations from the entire call stack up to but not including code in lit's
+#   implementation.
+# - RUN directive (plus line continuations).  Those are checked in detail in
+#   other tests.
+# - Each shell command.  Its trace appears in the stdout of the associated
+#   lit.run call or RUN directive.
+#
+# Execution traces when python code fails in some what are checked in
+# ./errors/trace-*.txt.
+
+# Check the execution trace details described above for config.prologue.
+#
+#  CHECK-NEXT:# executed config.prologue='[[PROLOGUE_FILE:.*/lit.prologue.py]]'
+#  CHECK-NEXT:# .---stdout from config.prologue='[[PROLOGUE_FILE]]'
+#  CHECK-NEXT:# | config.prologue writes to stdout
+#  CHECK-NEXT:# | # lit.run called from [[PROLOGUE_FILE]]:[[#]]
+#  CHECK-NEXT:# | true
+#  CHECK-NEXT:# | # executed command: true
+#  CHECK-NEXT:# | # lit.run called from [[PROLOGUE_FILE]]:[[#LOC_LN:]]
+#  CHECK-NEXT:# | echo config.prologue lit.run writes to stdout at line [[#LOC_LN]]
+#  CHECK-NEXT:# | # executed command: echo config.prologue lit.run writes to stdout at line [[#LOC_LN]]
+#  CHECK-NEXT:# | # .---command stdout------------
+#  CHECK-NEXT:# | # | config.prologue lit.run writes to stdout at line [[#LOC_LN]]
+#  CHECK-NEXT:# | # `-----------------------------
+#  CHECK-NEXT:# | # lit.run called from [[PROLOGUE_FILE]]:[[#LOC_LN:]]
+#  CHECK-NEXT:# | {{.*/python.*}} {{.*}}/write-to-stderr.py && echo at line [[#LOC_LN]]
+#  CHECK-NEXT:# | # executed command: {{.*/python.*}} {{.*}}/write-to-stderr.py
+#  CHECK-NEXT:# | # .---command stderr------------
+#  CHECK-NEXT:# | # | a line on stderr
+#  CHECK-NEXT:# | # `-----------------------------
+#  CHECK-NEXT:# | # executed command: echo at line [[#LOC_LN]]
+#  CHECK-NEXT:# | # .---command stdout------------
+#  CHECK-NEXT:# | # | at line [[#LOC_LN]]
+#  CHECK-NEXT:# | # `-----------------------------
+#  CHECK-NEXT:# | # lit.run called from [[PROLOGUE_FILE]]:[[#LOC_LN:]]
+#  CHECK-NEXT:# | {{.*/python.*}} {{.*}}/write-to-stdout-and-stderr.py && echo at line [[#LOC_LN]]
+#  CHECK-NEXT:# | # executed command: {{.*/python.*}} {{.*}}/write-to-stdout-and-stderr.py
+#  CHECK-NEXT:# | # .---command stdout------------
+#  CHECK-NEXT:# | # | a line on stdout
+#  CHECK-NEXT:# | # `-----------------------------
+#  CHECK-NEXT:# | # .---command stderr------------
+#  CHECK-NEXT:# | # | a line on stderr
+#  CHECK-NEXT:# | # `-----------------------------
+#  CHECK-NEXT:# | # executed command: echo at line [[#LOC_LN]]
+#  CHECK-NEXT:# | # .---command stdout------------
+#  CHECK-NEXT:# | # | at line [[#LOC_LN]]
+#  CHECK-NEXT:# | # `-----------------------------
+#  CHECK-NEXT:# | # lit.run called from [[PROLOGUE_FILE]]:[[#LOC_LN:]]
+#  CHECK-NEXT:# | #         called from [[PROLOGUE_FILE]]:[[#LOC_LN + 2]]
+#  CHECK-NEXT:# | #         called from [[PROLOGUE_FILE]]:[[#LOC_LN + 3]]
+#  CHECK-NEXT:# | echo config.prologue lit.run from func at line [[#LOC_LN]]
+#  CHECK-NEXT:# | # executed command: echo config.prologue lit.run from func at line [[#LOC_LN]]
+#  CHECK-NEXT:# | # .---command stdout------------
+#  CHECK-NEXT:# | # | config.prologue lit.run from func at line [[#LOC_LN]]
+#  CHECK-NEXT:# | # `-----------------------------
+#  CHECK-NEXT:# | # lit.run called from [[LOC_FILE:[^:]*]]:[[#LOC_LN:]]
+#  CHECK-NEXT:# | #         called from [[PROLOGUE_FILE]]:[[#]]
+#  CHECK-NEXT:# | echo hello world from lit.run at [[LOC_FILE]]:[[#LOC_LN]]
+#  CHECK-NEXT:# | # executed command: echo hello world from lit.run at [[LOC_FILE]]:[[#LOC_LN]]
+#  CHECK-NEXT:# | # .---command stdout------------
+#  CHECK-NEXT:# | # | hello world from lit.run at [[LOC_FILE]]:[[#LOC_LN]]
+#  CHECK-NEXT:# | # `-----------------------------
+#  CHECK-NEXT:# `-----------------------------
+#  CHECK-NEXT:# .---stderr from config.prologue='[[PROLOGUE_FILE]]'
+#  CHECK-NEXT:# | config.prologue writes to stderr
+#  CHECK-NEXT:# `-----------------------------
+
+# Repeat the above execution trace checks but now for a PYTHON directive block.
+#
+# Thus, this test checks the case when the first RUN or PYTHON directive is a
+# PYTHON directive.  ./prologue.txt checks the case when it's a RUN directive.
+#
+# PYTHON: import sys                                            #LN_PYTHON
+# PYTHON: print("PYTHON writes to stdout")                      #LN_PYTHON + 1
+# PYTHON: print("PYTHON writes to stderr", file=sys.stderr)     #LN_PYTHON + 2
+# PYTHON: lit.run("true")                                       #LN_PYTHON + 3
+# PYTHON: lit.run("echo PYTHON lit.run writes to stdout")       #LN_PYTHON + 4
+# PYTHON: lit.run("%{python} %S/write-to-stderr.py")            #LN_PYTHON + 5
+# PYTHON: lit.run("%{python} %S/write-to-stdout-and-stderr.py") #LN_PYTHON + 6
+# PYTHON: def localFnA():                                       #LN_PYTHON + 7
+# PYTHON:     lit.run("echo lit.run from func")                 #LN_PYTHON + 8
+# PYTHON: def localFnB():                                       #LN_PYTHON + 9
+# PYTHON:     localFnA()                                        #LN_PYTHON + 10
+# PYTHON: localFnB()                                            #LN_PYTHON + 11
+# PYTHON: exampleModule.helloWorldFromLitRun()                  #LN_PYTHON + 12
+#
+#  CHECK-NEXT:# executed '{{PYTHON}}:' directive from line [[#LN_PYTHON : @LINE - 14]] to [[# @LINE - 2]]
+#  CHECK-NEXT:# .---stdout from '{{PYTHON}}:' directive from line [[#LN_PYTHON]] to [[#LN_PYTHON + 12]]
+#  CHECK-NEXT:# | PYTHON writes to stdout
+#  CHECK-NEXT:# | # lit.run called from [[TEST_FILE]]:[[#LN_PYTHON + 3]]
+#  CHECK-NEXT:# | true
+#  CHECK-NEXT:# | # executed command: true
+#  CHECK-NEXT:# | # lit.run called from [[TEST_FILE]]:[[#LN_PYTHON + 4]]
+#  CHECK-NEXT:# | echo PYTHON lit.run writes to stdout
+#  CHECK-NEXT:# | # executed command: echo PYTHON lit.run writes to stdout
+#  CHECK-NEXT:# | # .---command stdout------------
+#  CHECK-NEXT:# | # | PYTHON lit.run writes to stdout
+#  CHECK-NEXT:# | # `-----------------------------
+#  CHECK-NEXT:# | # lit.run called from [[TEST_FILE]]:[[#LN_PYTHON + 5]]
+#  CHECK-NEXT:# | {{.*/python.*}} {{.*}}/write-to-stderr.py
+#  CHECK-NEXT:# | # executed command: {{.*/python.*}} {{.*}}/write-to-stderr.py
+#  CHECK-NEXT:# | # .---command stderr------------
+#  CHECK-NEXT:# | # | a line on stderr
+#  CHECK-NEXT:# | # `-----------------------------
+#  CHECK-NEXT:# | # lit.run called from [[TEST_FILE]]:[[#LN_PYTHON + 6]]
+#  CHECK-NEXT:# | {{.*/python.*}} {{.*}}/write-to-stdout-and-stderr.py
+#  CHECK-NEXT:# | # executed command: {{.*/python.*}} {{.*}}/write-to-stdout-and-stderr.py
+#  CHECK-NEXT:# | # .---command stdout------------
+#  CHECK-NEXT:# | # | a line on stdout
+#  CHECK-NEXT:# | # `-----------------------------
+#  CHECK-NEXT:# | # .---command stderr------------
+#  CHECK-NEXT:# | # | a line on stderr
+#  CHECK-NEXT:# | # `-----------------------------
+#  CHECK-NEXT:# | # lit.run called from [[TEST_FILE]]:[[#LN_PYTHON + 8]]
+#  CHECK-NEXT:# | #         called from [[TEST_FILE]]:[[#LN_PYTHON + 10]]
+#  CHECK-NEXT:# | #         called from [[TEST_FILE]]:[[#LN_PYTHON + 11]]
+#  CHECK-NEXT:# | echo lit.run from func
+#  CHECK-NEXT:# | # executed command: echo lit.run from func
+#  CHECK-NEXT:# | # .---command stdout------------
+#  CHECK-NEXT:# | # | lit.run from func
+#  CHECK-NEXT:# | # `-----------------------------
+#  CHECK-NEXT:# | # lit.run called from [[LOC_FILE:[^:]*]]:[[#LOC_LN:]]
+#  CHECK-NEXT:# | #         called from [[TEST_FILE]]:[[#LN_PYTHON + 12]]
+#  CHECK-NEXT:# | echo hello world from lit.run at [[LOC_FILE]]:[[#LOC_LN]]
+#  CHECK-NEXT:# | # executed command: echo hello world from lit.run at [[LOC_FILE]]:[[#LOC_LN]]
+#  CHECK-NEXT:# | # .---command stdout------------
+#  CHECK-NEXT:# | # | hello world from lit.run at [[LOC_FILE]]:[[#LOC_LN]]
+#  CHECK-NEXT:# | # `-----------------------------
+#  CHECK-NEXT:# `-----------------------------
+#  CHECK-NEXT:# .---stderr from '{{PYTHON}}:' directive from line [[#LN_PYTHON]] to [[#LN_PYTHON + 12]]
+#  CHECK-NEXT:# | PYTHON writes to stderr
+#  CHECK-NEXT:# `-----------------------------
+
+# This splits apart a PYTHON directive block.  That facilitates writing this
+# test, but it also checks that lit produces separate comments for separate
+# PYTHON directive blocks.
+#
+#         RUN: true
+#  CHECK-NEXT:# {{RUN}}: at line [[# @LINE - 1]]
+#  CHECK-NEXT:true
+#  CHECK-NEXT:# executed command: true
+
+# If there's only stdout from a PYTHON directive block, it should not print the
+# stderr header.  This is also checks the trace when the block contains only one
+# PYTHON directive.
+#
+# PYTHON: print("PYTHON writes only to stdout")
+#
+#  CHECK-NEXT:# executed '{{PYTHON}}:' directive at line [[#LN_PYTHON_STDOUT: @LINE - 2]]
+#  CHECK-NEXT:# .---stdout from '{{PYTHON}}:' directive at line [[#LN_PYTHON_STDOUT]]
+#  CHECK-NEXT:# | PYTHON writes only to stdout
+#  CHECK-NEXT:# `-----------------------------
+
+# Split PYTHON directive blocks so we can check the trace separately.
+#
+#         RUN: true
+#  CHECK-NEXT:# {{RUN}}: at line [[# @LINE - 1]]
+#  CHECK-NEXT:true
+#  CHECK-NEXT:# executed command: true
+
+# If there's only stderr from a PYTHON directive block, it should not print the
+# stdout header.
+#
+# PYTHON: print("PYTHON writes only to stderr", file=sys.stderr)
+#
+#  CHECK-NEXT:# executed '{{PYTHON}}:' directive at line [[#LN_PYTHON_STDERR: @LINE - 2]]
+#  CHECK-NEXT:# .---stderr from '{{PYTHON}}:' directive at line [[#LN_PYTHON_STDERR]]
+#  CHECK-NEXT:# | PYTHON writes only to stderr
+#  CHECK-NEXT:# `-----------------------------
+
+# Split PYTHON directive blocks so we can check the trace separately.
+#
+#         RUN: true
+#  CHECK-NEXT:# {{RUN}}: at line [[# @LINE - 1]]
+#  CHECK-NEXT:true
+#  CHECK-NEXT:# executed command: true
+
+# If there's no stdout or stderr, it shouldn't print either header.
+#
+# PYTHON: pass
+#
+#  CHECK-NEXT:# executed '{{PYTHON}}:' directive at line [[#LN_PYTHON_STDERR: @LINE - 2]]
+
+# Split PYTHON directive blocks so we can check the trace separately.
+#
+#         RUN: true
+#  CHECK-NEXT:# {{RUN}}: at line [[# @LINE - 1]]
+#  CHECK-NEXT:true
+#  CHECK-NEXT:# executed command: true
+
+# Check that reported line numbers are correct when a PYTHON directive block has
+# non-directive lines mixed in.
+#
+# PYTHON: lit.run("true") # LN_PYTHON_SEP_START
+
+# PYTHON: lit.run("true")
+# non-blank line
+# PYTHON: lit.run("true")
+#         non-blank line
+#
+# non-blank line
+# PYTHON: lit.run("true") # LN_PYTHON_SEP_END
+#
+#  CHECK-NEXT:# executed '{{PYTHON}}:' directive from line [[#LN_PYTHON_SEP_START: @LINE - 10]] to [[#LN_PYTHON_SEP_END: @LINE - 2]]
+#  CHECK-NEXT:# .---stdout from '{{PYTHON}}:' directive from line [[#LN_PYTHON_SEP_START]] to [[#LN_PYTHON_SEP_END]]
+#  CHECK-NEXT:# | # lit.run called from [[TEST_FILE]]:[[#LN_PYTHON_SEP_START]]
+#  CHECK-NEXT:# | true
+#  CHECK-NEXT:# | # executed command: true
+#  CHECK-NEXT:# | # lit.run called from [[TEST_FILE]]:[[#LN_PYTHON_SEP_START + 2]]
+#  CHECK-NEXT:# | true
+#  CHECK-NEXT:# | # executed command: true
+#  CHECK-NEXT:# | # lit.run called from [[TEST_FILE]]:[[#LN_PYTHON_SEP_START + 4]]
+#  CHECK-NEXT:# | true
+#  CHECK-NEXT:# | # executed command: true
+#  CHECK-NEXT:# | # lit.run called from [[TEST_FILE]]:[[#LN_PYTHON_SEP_START + 8]]
+#  CHECK-NEXT:# | true
+#  CHECK-NEXT:# | # executed command: true
+#  CHECK-NEXT:# `-----------------------------
+
+# End of "Command Output (stdout)".
+#
+# CHECK-EMPTY:
+#  CHECK-NEXT:--
+
+# CHECK:  Passed: 1
Index: llvm/utils/lit/tests/Inputs/shtest-python/write-to-stderr.py
===================================================================
--- /dev/null
+++ llvm/utils/lit/tests/Inputs/shtest-python/write-to-stderr.py
@@ -0,0 +1,6 @@
+#!/usr/bin/env python
+
+import sys
+
+sys.stderr.write("a line on stderr\n")
+sys.stderr.flush()
Index: llvm/utils/lit/tests/Inputs/shtest-python/write-to-stdout-and-stderr.py
===================================================================
--- /dev/null
+++ llvm/utils/lit/tests/Inputs/shtest-python/write-to-stdout-and-stderr.py
@@ -0,0 +1,9 @@
+#!/usr/bin/env python
+
+import sys
+
+sys.stdout.write("a line on stdout\n")
+sys.stdout.flush()
+
+sys.stderr.write("a line on stderr\n")
+sys.stderr.flush()
Index: llvm/utils/lit/tests/lit.cfg
===================================================================
--- llvm/utils/lit/tests/lit.cfg
+++ llvm/utils/lit/tests/lit.cfg
@@ -62,6 +62,8 @@
 # to ensure predictable test order, as it is often required for FileCheck
 # matches.
 config.substitutions.append(("%{inputs}", "Inputs"))
+config.substitutions.append(("%{inputs-absolute}",
+                             os.path.join(config.test_source_root, "Inputs")))
 config.substitutions.append(("%{lit}", "%{lit-no-order-opt} --order=lexical"))
 config.substitutions.append(
     (
Index: llvm/utils/lit/tests/per-test-coverage-by-lit-cfg.py
===================================================================
--- llvm/utils/lit/tests/per-test-coverage-by-lit-cfg.py
+++ llvm/utils/lit/tests/per-test-coverage-by-lit-cfg.py
@@ -18,7 +18,7 @@
 #      CHECK: export
 #      CHECK: LLVM_PROFILE_FILE=per-test-coverage-by-lit-cfg1.profraw
 #      CHECK: per-test-coverage-by-lit-cfg.py
-#      CHECK: {{RUN}}: at line 3
+#      CHECK: {{RUN}}: at line 4
 #      CHECK: export
-#      CHECK: LLVM_PROFILE_FILE=per-test-coverage-by-lit-cfg2.profraw
+#      CHECK: LLVM_PROFILE_FILE=per-test-coverage-by-lit-cfg3.profraw
 #      CHECK: per-test-coverage-by-lit-cfg.py
Index: llvm/utils/lit/tests/per-test-coverage.py
===================================================================
--- llvm/utils/lit/tests/per-test-coverage.py
+++ llvm/utils/lit/tests/per-test-coverage.py
@@ -18,7 +18,7 @@
 #      CHECK: export
 #      CHECK: LLVM_PROFILE_FILE=per-test-coverage1.profraw
 #      CHECK: per-test-coverage.py
-#      CHECK: {{RUN}}: at line 3
+#      CHECK: {{RUN}}: at line 4
 #      CHECK: export
-#      CHECK: LLVM_PROFILE_FILE=per-test-coverage2.profraw
+#      CHECK: LLVM_PROFILE_FILE=per-test-coverage3.profraw
 #      CHECK: per-test-coverage.py
Index: llvm/utils/lit/tests/shtest-keyword-parse-errors.py
===================================================================
--- llvm/utils/lit/tests/shtest-keyword-parse-errors.py
+++ llvm/utils/lit/tests/shtest-keyword-parse-errors.py
@@ -6,7 +6,7 @@
 # CHECK: Testing: 3 tests
 
 # CHECK-LABEL: UNRESOLVED: shtest-keyword-parse-errors :: empty.txt
-# CHECK:       {{^}}Test has no 'RUN:' line{{$}}
+# CHECK:       {{^}}Test has no 'RUN:' line or 'PYTHON:' line{{$}}
 
 # CHECK-LABEL: UNRESOLVED: shtest-keyword-parse-errors :: multiple-allow-retries.txt
 # CHECK:       {{^}}Test has more than one ALLOW_RETRIES lines{{$}}
Index: llvm/utils/lit/tests/shtest-python.py
===================================================================
--- /dev/null
+++ llvm/utils/lit/tests/shtest-python.py
@@ -0,0 +1,50 @@
+# We're using PYTHON to help us write tests for PYTHON.
+
+# RUN: echo "-- Available Tests --" > %t.tests.actual.txt
+
+# DEFINE: %{my-inputs} = %{inputs-absolute}/shtest-python
+
+# PYTHON: def runTest(test, litPre="", prologue=None):
+# PYTHON:     litCmd = litPre + " %{lit} -va"
+# PYTHON:     if prologue:
+# PYTHON:         litCmd += " -DaddPrologue=" + prologue
+# PYTHON:     testAbs = "%{my-inputs}/" + test
+# PYTHON:     lit.run(f"""
+# PYTHON:         {litCmd} {testAbs} 2>&1 |
+# PYTHON:           FileCheck {testAbs} -strict-whitespace -match-full-lines
+# PYTHON:                               -dump-input-filter=all -vv -color
+# PYTHON:                               -DTEST_FILE='{testAbs}' &&
+# PYTHON:         echo "  shtest-python :: {test}" >> %t.tests.actual.txt
+# PYTHON:     """)
+# PYTHON: def runPass(test, prologue=None):
+# PYTHON:     runTest(test, prologue=prologue)
+# PYTHON: def runFail(test, prologue=None):
+# PYTHON:     runTest(test, litPre="not", prologue=prologue)
+
+# PYTHON: runFail("errors/incomplete-python.txt")
+# PYTHON: runFail("errors/inconsistent-indent-2-lines.txt")
+# PYTHON: runFail("errors/inconsistent-indent-3-lines.txt")
+# PYTHON: runFail("errors/internal-api-inaccessible.txt")
+# PYTHON: runFail("errors/python-syntax-error-1-line.txt")
+# PYTHON: runFail("errors/python-syntax-error-2-lines.txt")
+# PYTHON: runFail("errors/python-syntax-error-3-lines.txt")
+# PYTHON: runFail("errors/python-syntax-error-leading-indent.txt")
+# PYTHON: runFail("errors/trace-compile-error.txt")
+# PYTHON: runFail("errors/trace-exception.txt")
+# PYTHON: runFail("errors/trace-prologue-exception.txt",
+# PYTHON:         prologue="errors/lit.prologue.py")
+# PYTHON: runFail("errors/trace-run-compile-error.txt")
+
+# PYTHON: runPass("file.txt")
+# PYTHON: runPass("has.txt")
+# PYTHON: runPass("import.txt")
+# PYTHON: runPass("no-shell-commands.txt")
+# PYTHON: runPass("prologue.txt", prologue="lit.prologue.py")
+# PYTHON: runPass("shell-affects-python.txt")
+# PYTHON: runPass("substs-affects-python.txt")
+# PYTHON: runPass("trace.txt", prologue="lit.prologue.py")
+
+# Make sure we didn't forget to run something.
+#
+# RUN: %{lit} --show-tests %{my-inputs} > %t.tests.expected.txt
+# RUN: diff -u -w %t.tests.expected.txt %t.tests.actual.txt