[llvm-lit] Adding -e option to lit's built-in cat command

llvm/utils/lit/lit/builtin_commands/cat.py

@@ -1,40 +1,59 @@
 import getopt
 import sys
+from dataclasses import dataclass
 
 try:
     from StringIO import StringIO
 except ImportError:
     from io import StringIO
 
 
-def convertToCaretAndMNotation(data):
+# This dataclass defines all currently supported options for cat
+@dataclass
+class Options:
+    # Options: -e. True if newlines should be displayed with a '$'
+    show_ends: bool
+    # Options: -v, -e. True if text should be converted to ^ and M- notation
+    show_nonprinting: bool
+
+
+def convertTextNotation(data, options):
+    assert options.show_ends or options.show_nonprinting
+
     newdata = StringIO()
     if isinstance(data, str):
         data = bytearray(data.encode())
 
     for intval in data:
-        if intval == 9 or intval == 10:
+        if intval == 10 and options.show_ends:
+            newdata.write("$")
             newdata.write(chr(intval))
             continue
-        if intval > 127:
-            intval = intval - 128
-            newdata.write("M-")
-        if intval < 32:
-            newdata.write("^")
-            newdata.write(chr(intval + 64))
-        elif intval == 127:
-            newdata.write("^?")
-        else:
-            newdata.write(chr(intval))
+        if options.show_nonprinting:
+            if intval == 9 or intval == 10:
+                newdata.write(chr(intval))
+                continue
+            if intval > 127:
+                intval = intval - 128
+                newdata.write("M-")
+            if intval < 32:
+                newdata.write("^")
+                newdata.write(chr(intval + 64))
+                continue
+            elif intval == 127:
+                newdata.write("^?")
+                continue
+        newdata.write(chr(intval))
 
     return newdata.getvalue().encode()
 
 
 def main(argv):
     arguments = argv[1:]
-    short_options = "v"
-    long_options = ["show-nonprinting"]
-    show_nonprinting = False
+    short_options = "eEv"
+    long_options = ["show-ends", "show-nonprinting"]
+    enabled_options = Options(show_ends=False, show_nonprinting=False)
+    convert_text = False
 
     try:
         options, filenames = getopt.gnu_getopt(arguments, short_options, long_options)
@@ -43,8 +62,12 @@ def main(argv):
         sys.exit(1)
 
     for option, value in options:
-        if option == "-v" or option == "--show-nonprinting":
-            show_nonprinting = True
+        if option == "-v" or option == "--show-nonprinting" or option == "-e":
+            enabled_options.show_nonprinting = True
+            convert_text = True
+        if option == "-E" or option == "--show-ends" or option == "-e":
+            enabled_options.show_ends = True
+            convert_text = True
 
     writer = getattr(sys.stdout, "buffer", None)
     if writer is None:
@@ -69,8 +92,8 @@ def main(argv):
                 fileToCat = open(filename, "rb")
                 contents = fileToCat.read()
 
-            if show_nonprinting:
-                contents = convertToCaretAndMNotation(contents)
+            if convert_text:
+                contents = convertTextNotation(contents, enabled_options)
             elif is_text:
                 contents = contents.encode()
             writer.write(contents)