[llvm] [lit] Use provided streams in builtin diff (PR #204869)

Prasoon Kumar via llvm-commits llvm-commits at lists.llvm.org
Mon Jun 29 18:52:55 PDT 2026


https://github.com/prasoon054 updated https://github.com/llvm/llvm-project/pull/204869

>From e6f621ed7264e5bd268c68e14a085123f66a2b56 Mon Sep 17 00:00:00 2001
From: Prasoon Kumar <prasoonkumar054 at gmail.com>
Date: Fri, 19 Jun 2026 22:56:50 +0530
Subject: [PATCH] [lit] Use provided streams in builtin diff

We want to move the diff builtin to run in-process inside the lit
worker, instead of spawning a subprocess. The current implementation
talks to sys.stdin / stdout / stderr directly, so it can't be called with
different streams.

To fix this, pull diff's logic into run(argv, stdin, stdout, stderr,
cwd), which takes streams as arguments instead of reaching for sys.std*.
main() now just calls run() with the real process streams, so the
spawned-script path is unchanged.

This also makes the 'import util' dual-mode: lit.util when diff is
imported as part of the lit package, falling back to flat util for the
spawned script.
---
 llvm/utils/lit/lit/builtin_commands/diff.py | 133 +++++++++++++-------
 1 file changed, 91 insertions(+), 42 deletions(-)

diff --git a/llvm/utils/lit/lit/builtin_commands/diff.py b/llvm/utils/lit/lit/builtin_commands/diff.py
index 5eab10fe17165..579a2649bc0af 100644
--- a/llvm/utils/lit/lit/builtin_commands/diff.py
+++ b/llvm/utils/lit/lit/builtin_commands/diff.py
@@ -1,13 +1,26 @@
 import difflib
 import functools
 import getopt
-import io
 import locale
 import os
 import re
 import sys
 
-import util
+# diff.py runs in two modes during the in-process migration:
+#   - In-process (default): imported as 'lit.builtin_commands.diff', so the
+#     'lit' package is already resolvable => lit.util works
+#   - Spawned fallback: run standalone via 'sys.executable diff.py', with
+#     PYTHONPATH set to the 'lit/' directory itself, not its parent. No entry
+#     on sys.path is the parent of 'lit/', so 'lit' can't be imported as package =>
+#     'lit.util' raises ImportError. util.py sits directly in that PYTHONPATH
+#     dir, so bare 'import util' resolves instead.
+# Both paths load the same module; the fallback handles the path difference.
+# TODO: Collapse to a standard 'import lit.util' once standalone spawning is removed.
+
+try:
+    import lit.util as util
+except ImportError:
+    import util
 
 
 class DiffFlags:
@@ -48,29 +61,27 @@ def getDirTree(path, basedir=""):
         return path, sorted(child_trees)
 
 
-def compareTwoFiles(flags, filepaths):
+def compareTwoFiles(flags, filepaths, stdin, stdout):
     filelines = []
     for file in filepaths:
         if file == "-":
-            stdin_fileno = sys.stdin.fileno()
-            with os.fdopen(os.dup(stdin_fileno), "rb") as stdin_bin:
-                filelines.append(stdin_bin.readlines())
+            filelines.append(stdin.readlines())
         else:
             with open(file, "rb") as file_bin:
                 filelines.append(file_bin.readlines())
 
     try:
         return compareTwoTextFiles(
-            flags, filepaths, filelines, locale.getpreferredencoding(False)
+            flags, filepaths, filelines, locale.getpreferredencoding(False), stdout
         )
     except UnicodeDecodeError:
         try:
-            return compareTwoTextFiles(flags, filepaths, filelines, "utf-8")
+            return compareTwoTextFiles(flags, filepaths, filelines, "utf-8", stdout)
         except:
-            return compareTwoBinaryFiles(flags, filepaths, filelines)
+            return compareTwoBinaryFiles(flags, filepaths, filelines, stdout)
 
 
-def compareTwoBinaryFiles(flags, filepaths, filelines):
+def compareTwoBinaryFiles(flags, filepaths, filelines, stdout):
     exitCode = 0
     diffs = difflib.diff_bytes(
         difflib.unified_diff,
@@ -82,12 +93,16 @@ def compareTwoBinaryFiles(flags, filepaths, filelines):
     )
 
     for diff in diffs:
-        sys.stdout.write(diff.decode(errors="backslashreplace"))
+        stdout.write(
+            diff.decode(errors="backslashreplace").encode(
+                locale.getpreferredencoding(False)
+            )
+        )
         exitCode = 1
     return exitCode
 
 
-def compareTwoTextFiles(flags, filepaths, filelines_bin, encoding):
+def compareTwoTextFiles(flags, filepaths, filelines_bin, encoding, stdout):
     filelines = []
     for lines_bin in filelines_bin:
         lines = []
@@ -129,32 +144,40 @@ def compose2(f, g):
         filepaths[1],
         n=flags.num_context_lines,
     ):
-        sys.stdout.write(diff)
+        stdout.write(diff.encode(encoding))
         exitCode = 1
     return exitCode
 
 
-def printDirVsFile(dir_path, file_path):
+def printDirVsFile(dir_path, file_path, stdout):
     if os.path.getsize(file_path):
         msg = "File %s is a directory while file %s is a regular file"
     else:
         msg = "File %s is a directory while file %s is a regular empty file"
-    sys.stdout.write(msg % (dir_path, file_path) + "\n")
+    stdout.write(
+        (msg % (dir_path, file_path) + "\n").encode(locale.getpreferredencoding(False))
+    )
 
 
-def printFileVsDir(file_path, dir_path):
+def printFileVsDir(file_path, dir_path, stdout):
     if os.path.getsize(file_path):
         msg = "File %s is a regular file while file %s is a directory"
     else:
         msg = "File %s is a regular empty file while file %s is a directory"
-    sys.stdout.write(msg % (file_path, dir_path) + "\n")
+    stdout.write(
+        (msg % (file_path, dir_path) + "\n").encode(locale.getpreferredencoding(False))
+    )
 
 
-def printOnlyIn(basedir, path, name):
-    sys.stdout.write("Only in %s: %s\n" % (os.path.join(basedir, path), name))
+def printOnlyIn(basedir, path, name, stdout):
+    stdout.write(
+        ("Only in %s: %s\n" % (os.path.join(basedir, path), name)).encode(
+            locale.getpreferredencoding(False)
+        )
+    )
 
 
-def compareDirTrees(flags, dir_trees, base_paths=["", ""]):
+def compareDirTrees(flags, dir_trees, stdin, stdout, base_paths=["", ""]):
     # Dirnames of the trees are not checked, it's caller's responsibility,
     # as top-level dirnames are always different. Base paths are important
     # for doing os.walk, but we don't put it into tree's dirname in order
@@ -170,12 +193,15 @@ def compareDirTrees(flags, dir_trees, base_paths=["", ""]):
                 os.path.join(left_base, left_tree[0]),
                 os.path.join(right_base, right_tree[0]),
             ],
+            stdin,
+            stdout,
         )
 
     if left_tree[1] is None and right_tree[1] is not None:
         printFileVsDir(
             os.path.join(left_base, left_tree[0]),
             os.path.join(right_base, right_tree[0]),
+            stdout,
         )
         return 1
 
@@ -183,6 +209,7 @@ def compareDirTrees(flags, dir_trees, base_paths=["", ""]):
         printDirVsFile(
             os.path.join(left_base, left_tree[0]),
             os.path.join(right_base, right_tree[0]),
+            stdout,
         )
         return 1
 
@@ -195,16 +222,18 @@ def compareDirTrees(flags, dir_trees, base_paths=["", ""]):
         # Names are sorted in getDirTree, rely on that order.
         if left_names[l] < right_names[r]:
             exitCode = 1
-            printOnlyIn(left_base, left_tree[0], left_names[l])
+            printOnlyIn(left_base, left_tree[0], left_names[l], stdout)
             l += 1
         elif left_names[l] > right_names[r]:
             exitCode = 1
-            printOnlyIn(right_base, right_tree[0], right_names[r])
+            printOnlyIn(right_base, right_tree[0], right_names[r], stdout)
             r += 1
         else:
             exitCode |= compareDirTrees(
                 flags,
                 [left_tree[1][l], right_tree[1][r]],
+                stdin,
+                stdout,
                 [
                     os.path.join(left_base, left_tree[0]),
                     os.path.join(right_base, right_tree[0]),
@@ -216,28 +245,44 @@ def compareDirTrees(flags, dir_trees, base_paths=["", ""]):
     # At least one of the trees has ended. Report names from the other tree.
     while l < len(left_names):
         exitCode = 1
-        printOnlyIn(left_base, left_tree[0], left_names[l])
+        printOnlyIn(left_base, left_tree[0], left_names[l], stdout)
         l += 1
     while r < len(right_names):
         exitCode = 1
-        printOnlyIn(right_base, right_tree[0], right_names[r])
+        printOnlyIn(right_base, right_tree[0], right_names[r], stdout)
         r += 1
     return exitCode
 
 
-def main(argv):
-    if sys.platform == "win32":
-        sys.stdout = io.TextIOWrapper(sys.stdout.buffer, newline="\n")
+def run(argv, stdin, stdout, stderr, cwd):
+    """In-process diff.
 
+    Writes bytes to stdout, returns an exit code, and never calls sys.exit.
+    This makes it safe to run inside lit worker as well as from the
+    standalone main below.
+
+    Args:
+        argv: Command-line arguments. The first element is the command name,
+            followed by options and the two files or directories to compare.
+        stdin: Binary input stream, used if a file operand is '-'.
+        stdout: Binary output stream for the diff output.
+        stderr: Binary error stream for error messages.
+        cwd: The shell's current working directory, used to resolve relative
+            file paths.
+
+    Returns:
+        An integer representing the exit code (0 for no differences, 1 if
+        differences were found or an error occurred).
+    """
     args = argv[1:]
     try:
         opts, args = getopt.gnu_getopt(args, "wbuI:U:r", ["strip-trailing-cr"])
     except getopt.GetoptError as err:
-        sys.stderr.write("Unsupported: 'diff': %s\n" % str(err))
-        sys.exit(1)
+        stderr.write(b"Unsupported: 'diff': %s\n" % str(err).encode())
+        return 1
 
     flags = DiffFlags()
-    filelines, filepaths, dir_trees = ([] for i in range(3))
+    filepaths, dir_trees = [], []
     for o, a in opts:
         if o == "-w":
             flags.ignore_all_space = True
@@ -250,10 +295,10 @@ def main(argv):
             try:
                 flags.num_context_lines = int(a)
                 if flags.num_context_lines < 0:
-                    raise ValueException
+                    raise ValueError
             except:
-                sys.stderr.write("Error: invalid '-U' argument: {}\n".format(a))
-                sys.exit(1)
+                stderr.write(b"Error: invalid '-U' argument: %s\n" % a.encode())
+                return 1
         elif o == "-I":
             flags.ignore_matching_lines = True
             flags.ignore_matching_lines_regex = a
@@ -265,34 +310,38 @@ def main(argv):
             assert False, "unhandled option"
 
     if len(args) != 2:
-        sys.stderr.write("Error: missing or extra operand\n")
-        sys.exit(1)
+        stderr.write(b"Error: missing or extra operand\n")
+        return 1
 
     exitCode = 0
     try:
         for file in args:
             if file != "-" and not os.path.isabs(file):
-                file = util.abs_path_preserve_drive(file)
+                file = util.abs_path_preserve_drive(os.path.join(cwd, file))
 
             if flags.recursive_diff:
                 if file == "-":
-                    sys.stderr.write("Error: cannot recursively compare '-'\n")
-                    sys.exit(1)
+                    stderr.write(b"Error: cannot recursively compare '-'\n")
+                    return 1
                 dir_trees.append(getDirTree(file))
             else:
                 filepaths.append(file)
 
         if not flags.recursive_diff:
-            exitCode = compareTwoFiles(flags, filepaths)
+            exitCode = compareTwoFiles(flags, filepaths, stdin, stdout)
         else:
-            exitCode = compareDirTrees(flags, dir_trees)
+            exitCode = compareDirTrees(flags, dir_trees, stdin, stdout)
 
     except IOError as err:
-        sys.stderr.write("Error: 'diff' command failed, %s\n" % str(err))
+        stderr.write(b"Error: 'diff' command failed, %s\n" % str(err).encode())
         exitCode = 1
 
-    sys.exit(exitCode)
+    return exitCode
+
 
+def main(argv):
+    out = getattr(sys.stdout, "buffer", sys.stdout)
+    sys.exit(run(argv, sys.stdin.buffer, out, sys.stderr.buffer, os.getcwd()))
 
 if __name__ == "__main__":
     main(sys.argv)



More information about the llvm-commits mailing list