1 files changed, 112 insertions, 0 deletions
diff --git a/rainbow_parenthesis/__init__.py b/rainbow_parenthesis/__init__.py
new file mode 100644
index 0000000..ad8eb51
--- /dev/null
+++ b/rainbow_parenthesis/__init__.py
@@ -0,0 +1,112 @@
+"""Simple script which adds matching rainbow colors to data read on stdin."""
+
+import io
+from typing import Literal
+from dataclasses import dataclass
+
+
+@dataclass
+class Stackpointer():
+    """A "fancy" integer which implements pre-and post decrement."""
+
+    depth: int = 0
+
+    def __call__(self, c: int):
+        """
+        Update and return current state.
+
+        :param c:
+            An integer.
+            Positive integers increment the value after returning it,
+            negative integers decrement the value before returning it,
+            0 returns the value as is.
+        """
+        ret: int
+        if c > 0:
+            ret = self.depth
+            self.depth += 1
+        elif c < 0:
+            self.depth -= 1
+            ret = self.depth
+        else:
+            ret = self.depth
+        return ret
+
+
+@dataclass
+class Colored:
+    """
+    Tag an item with a color "depth".
+
+    Depth is an arbitarary (positive) integer, from 0 and
+    incrementing by up. Each distinct value should correspond to a
+    color. The colors may repeat.
+    """
+
+    depth: int
+    item: str
+
+
+def color(depth: int, c: str) -> Colored:
+    """Write a highlighted string."""
+    return Colored(depth, c)
+
+
+def colorize(strm: io.TextIOBase) -> list[str | Colored]:
+    """
+    Colorize a given string.
+
+    :param strm:
+        Text stream to get contents from.
+
+        Use ``io.StringIO`` if you want to pass a string.
+    :returns:
+        A list where each item is either a plain string, or a
+        ``Colored`` object.
+    """
+    in_string: Literal[False] | Literal['"'] | Literal["'"] = False
+
+    paren = 0
+    brace = 0
+    brack = 0
+
+    depth = Stackpointer()
+
+    out: list[str | Colored] = []
+    while c := strm.read(1):
+        match c:
+            case '(':
+                out.append(color(depth(1), c))
+                paren += 1
+            case ')':
+                out.append(color(depth(-1), c))
+                paren -= 1
+            case '[':
+                out.append(color(depth(1), c))
+                brack += 1
+            case ']':
+                out.append(color(depth(-1), c))
+                brack -= 1
+            case '{':
+                out.append(color(depth(1), c))
+                brace += 1
+            case '}':
+                out.append(color(depth(-1), c))
+                brace -= 1
+            case "'":
+                if in_string == "'":
+                    in_string = False
+                else:
+                    in_string = "'"
+                out.append(c)
+            case '"':
+                if in_string == '"':
+                    in_string = False
+                else:
+                    in_string = '"'
+                out.append(c)
+            case '\\':
+                out.append(strm.read(1))
+            case c:
+                out.append(c)
+    return out