1 files changed, 69 insertions, 0 deletions
diff --git a/scripts/utility.py b/scripts/utility.py
new file mode 100644
index 0000000..4d59a1b
--- /dev/null
+++ b/scripts/utility.py
@@ -0,0 +1,69 @@
+"""
+    Utility functions for test scripts
+    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+    :copyright: Copyright 2006-2022 by the Pygments team, see AUTHORS.
+    :license: BSD, see LICENSE for details.
+"""
+
+import os
+import os.path
+
+
+def unpack_output_file(path):
+    """
+    Unpack an output file into objects contining the line number, the text,
+    and the token name. The output file can be either a ``.output`` file
+    containing a token stream, or a ``.txt`` with input and tokens.
+    """
+    from collections import namedtuple
+    entry = namedtuple('OutputEntry', ['text', 'token', 'linenumber'])
+
+    skip_until_tokens = path.endswith('.txt')
+
+    for linenumber, line in enumerate(open(path).readlines()):
+        line = line.strip()
+        if not line:
+            continue
+
+        if skip_until_tokens:
+            if line != '---tokens---':
+                continue
+            else:
+                skip_until_tokens = False
+                continue
+
+        # Line can start with ' or ", so let's check which one it is
+        # and find the matching one
+        quotation_start = 0
+        quotation_end = line.rfind(line[0])
+        text = line[quotation_start+1:quotation_end]
+        token = line.split()[-1]
+        text = text.replace('\\n', '\n')
+        text = text.replace('\\t', '\t')
+        yield entry(text, token, linenumber + 1)
+
+
+def process_output_files(root_directory, callback):
+    """
+    Process all output (i.e. .output and .txt files for snippets) files
+    in a directory tree using the provided callback.
+    The callback should return ``True`` in case of success, ``False``
+    otherwise.
+
+    The function returns the number of files for which the callback returned
+    ``False``.
+    """
+    errors = 0
+    for dir, _, files in os.walk(root_directory):
+        for file in files:
+            _, ext = os.path.splitext(file)
+
+            if ext not in {'.txt', '.output'}:
+                continue
+
+            path = os.path.join(dir, file)
+            if not callback(path):
+                errors += 1
+
+    return errors