Move project to a proper module layout.

5 files changed, 187 insertions, 122 deletions

diff --git a/Makefile b/Makefile
new file mode 100644
index 0000000..595a7c2
--- /dev/null
+++ b/Makefile
@@ -0,0 +1,5 @@
+.PHONY: check
+
+check:
+	flake8
+	mypy -p rainbow_parenthesis
diff --git a/main.py b/main.py
deleted file mode 100755
index 406549d..0000000
--- a/main.py
+++ /dev/null
@@ -1,122 +0,0 @@
-#!/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 = 31
-    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()
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
diff --git a/rainbow_parenthesis/__main__.py b/rainbow_parenthesis/__main__.py
new file mode 100644
index 0000000..251731a
--- /dev/null
+++ b/rainbow_parenthesis/__main__.py
@@ -0,0 +1,23 @@
+"""
+Entry point for rainbow parenthesis.
+
+Reads a string from stdin, and outputs it to stdout with all
+parenthesis prettily colored.
+"""
+
+from . import colorize, Colored
+from . import term
+import argparse
+
+parser = argparse.ArgumentParser(prog='rainbow')
+parser.add_argument('input', type=argparse.FileType('r'),
+                    nargs='?', default='-')
+args = parser.parse_args()
+
+
+for item in colorize(args.input):
+    match item:
+        case Colored():
+            print(term.colorize(item), end='')
+        case s:
+            print(s, end='')
diff --git a/rainbow_parenthesis/term.py b/rainbow_parenthesis/term.py
new file mode 100644
index 0000000..c712e74
--- /dev/null
+++ b/rainbow_parenthesis/term.py
@@ -0,0 +1,47 @@
+"""
+Output engine for rainbow parenthesis.
+
+Attaches terminal escape codes for use on UNIX-like systems.
+"""
+
+from enum import auto, Enum
+# from dataclasses import dataclass, field
+
+from . import Colored
+
+
+CSI = '\033['
+
+
+class ANSIColor(Enum):
+    """Known CSI colors."""
+
+    # BLACK = 30
+    RED = 31
+    GREEN = auto()
+    YELLOW = auto()
+    BLUE = auto()
+    MAGENTA = auto()
+    CYAN = auto()
+
+    @classmethod
+    def get(cls, depth: int) -> int:
+        """Get color code for the given depth."""
+        return list(cls)[depth % len(list(cls))].value
+
+
+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'
+
+
+def colorize(c: Colored) -> str:
+    """Return contents of colored, surrounded by color escapes."""
+    out = SGR(1, ANSIColor.get(c.depth))
+    out += c.item
+    out += SGR()
+    return out