aboutsummaryrefslogtreecommitdiff
path: root/rainbow_parenthesis/__init__.py
blob: ad8eb51f071ed51cba68f37130fd2f5321e465c3 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
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