Add code.

2 files changed, 128 insertions, 0 deletions

diff --git a/README.md b/README.md
new file mode 100644
index 0000000..bebbbd3
--- /dev/null
+++ b/README.md
@@ -0,0 +1,6 @@
+Rainbow Parenthesis
+===================
+
+A simple script which highlights parenthesis on the data given on stdin.
+
+    cat main.py | ./main.py
diff --git a/main.py b/main.py
new file mode 100755
index 0000000..d1ef47f
--- /dev/null
+++ b/main.py
@@ -0,0 +1,122 @@
+#!/usr/bin/env python3
+
+"""Simple script which adds matching rainbow colors to data read on stdin."""
+
+import sys
+from typing import Literal
+from enum import auto, Enum
+from dataclasses import dataclass, field
+
+CSI = '\033['
+
+
+def SGR(*xs):
+    """
+    Build a CSI escape sequence.
+
+    https://en.wikipedia.org/wiki/ANSI_escape_code#CSI_(Control_Sequence_Introducer)_sequences
+    """
+    return CSI + ';'.join(str(x) for x in xs) + 'm'
+
+
+class Color(Enum):
+    """Known CSI colors."""
+
+    BLACK = 30
+    RED = auto()
+    GREEN = auto()
+    YELLOW = auto()
+    BLUE = auto()
+    MAGENTA = auto()
+    CYAN = auto()
+
+
+@dataclass
+class ColorGenerator:
+    """
+    Handles colors by "depth".
+
+    :param colors:
+        List of colors to use.
+    :param _idx:
+        Current index.
+    """
+
+    colors: list[Color] = field(default_factory=lambda: list(Color))
+    _idx: int = 0
+
+    def __call__(self, c: int):
+        """
+        Update the current color, and return it.
+
+        :param c:
+            Any positive or negative number changes the interval value
+            by 1 or -1. 0 simply returns the current value.
+        """
+        if c > 0:
+            ret = self.colors[self._idx % len(self.colors)]
+            self._idx += 1
+        elif c < 0:
+            self._idx -= 1
+            ret = self.colors[self._idx % len(self.colors)]
+        else:
+            ret = self.colors[self._idx % len(self.colors)]
+        return ret
+
+
+def color(color: Color, c: str):
+    """Write a highlighted string."""
+    sys.stdout.write(SGR(1, color.value))
+    sys.stdout.write(c)
+    sys.stdout.write(SGR())
+
+
+def __main() -> None:
+    in_string: Literal[False] | Literal['"'] | Literal["'"] = False
+
+    paren = 0
+    brace = 0
+    brack = 0
+
+    col = ColorGenerator()
+
+    while c := sys.stdin.read(1):
+        match c:
+            case '(':
+                color(col(1), c)
+                paren += 1
+            case ')':
+                color(col(-1), c)
+                paren -= 1
+            case '[':
+                color(col(1), c)
+                brack += 1
+            case ']':
+                color(col(-1), c)
+                brack -= 1
+            case '{':
+                color(col(1), c)
+                brace += 1
+            case '}':
+                color(col(-1), c)
+                brace -= 1
+            case "'":
+                if in_string == "'":
+                    in_string = False
+                else:
+                    in_string = "'"
+                sys.stdout.write(c)
+            case '"':
+                if in_string == '"':
+                    in_string = False
+                else:
+                    in_string = '"'
+                sys.stdout.write(c)
+            case '\\':
+                sys.stdout.write(sys.stdin.read(1))
+            case c:
+                sys.stdout.write(c)
+
+
+if __name__ == '__main__':
+    __main()