Switch to using a generator.

2 files changed, 38 insertions, 29 deletions

diff --git a/rainbow_parenthesis/__init__.py b/rainbow_parenthesis/__init__.py
index ad8eb51..4be7d80 100644
--- a/rainbow_parenthesis/__init__.py
+++ b/rainbow_parenthesis/__init__.py
@@ -1,7 +1,12 @@
-"""Simple script which adds matching rainbow colors to data read on stdin."""
+"""
+Colorize matching parenthesis in a string.
+
+The "primary" entry point is ``colorize``, and ``Colored`` the
+"primary" data structure.
+"""
 
 import io
-from typing import Literal
+from typing import Literal, Generator
 from dataclasses import dataclass
 
 
@@ -39,8 +44,7 @@ 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.
+    incrementing by up. Should map to colors.
     """
 
     depth: int
@@ -52,7 +56,7 @@ def color(depth: int, c: str) -> Colored:
     return Colored(depth, c)
 
 
-def colorize(strm: io.TextIOBase) -> list[str | Colored]:
+def colorize(strm: io.TextIOBase) -> Generator[str | Colored, None, None]:
     """
     Colorize a given string.
 
@@ -61,8 +65,9 @@ def colorize(strm: io.TextIOBase) -> list[str | Colored]:
 
         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.
+        Yields a stream of tokens, where each token is either a
+        literal string, laternatively, it's a ``Colored``
+        object, which attaches a color number to the string.
     """
     in_string: Literal[False] | Literal['"'] | Literal["'"] = False
 
@@ -72,41 +77,39 @@ def colorize(strm: io.TextIOBase) -> list[str | Colored]:
 
     depth = Stackpointer()
 
-    out: list[str | Colored] = []
     while c := strm.read(1):
         match c:
             case '(':
-                out.append(color(depth(1), c))
+                yield color(depth(1), c)
                 paren += 1
             case ')':
-                out.append(color(depth(-1), c))
+                yield color(depth(-1), c)
                 paren -= 1
             case '[':
-                out.append(color(depth(1), c))
+                yield color(depth(1), c)
                 brack += 1
             case ']':
-                out.append(color(depth(-1), c))
+                yield color(depth(-1), c)
                 brack -= 1
             case '{':
-                out.append(color(depth(1), c))
+                yield color(depth(1), c)
                 brace += 1
             case '}':
-                out.append(color(depth(-1), c))
+                yield color(depth(-1), c)
                 brace -= 1
             case "'":
                 if in_string == "'":
                     in_string = False
                 else:
                     in_string = "'"
-                out.append(c)
+                yield c
             case '"':
                 if in_string == '"':
                     in_string = False
                 else:
                     in_string = '"'
-                out.append(c)
+                yield c
             case '\\':
-                out.append(strm.read(1))
+                yield strm.read(1)
             case c:
-                out.append(c)
-    return out
+                yield c
diff --git a/rainbow_parenthesis/__main__.py b/rainbow_parenthesis/__main__.py
index 251731a..51473c9 100644
--- a/rainbow_parenthesis/__main__.py
+++ b/rainbow_parenthesis/__main__.py
@@ -9,15 +9,21 @@ 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()
 
+def main():
+    """Interactive entry point for the program."""
+    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='')
+    for item in colorize(args.input):
+        match item:
+            case Colored():
+                print(term.colorize(item), end='')
+            case s:
+                print(s, end='')
+
+
+if __name__ == '__main__':
+    main()