aboutsummaryrefslogtreecommitdiff
path: root/rainbow_parenthesis/__init__.py
blob: e665bdb81d91d672b67a7191ba14eb7961a7da4f (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
113
114
115
116
117
118
119
120
121
122
123
124
125
"""
Colorize matching parenthesis in a string.

The "primary" entry point is ``colorize``, and ``Colored`` the
"primary" data structure.
"""

import io
from typing import Literal, Generator
from dataclasses import dataclass

__version__ = "0.2"


@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. Should map to colors.
    """

    depth: int
    item: str


def color(depth: int, c: str) -> Colored:
    """Write a highlighted string."""
    return Colored(depth, c)


def colorize(strm: io.TextIOBase) -> Generator[str | Colored, None, None]:
    """
    Colorize a given string.

    :param strm:
        Text stream to get contents from.

        Use ``io.StringIO`` if you want to pass a string.
    :returns:
        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

    paren = 0
    brace = 0
    brack = 0

    depth = Stackpointer()

    while c := strm.read(1):
        match [in_string, c]:
            case ['"', '"']:
                # The string is ending
                in_string = False
                yield c
            case ["'", "'"]:
                in_string = False
                # The string is ending
                yield c
            case [False, _]:
                # We are not in a string, emit the next character
                # colored if it's a parenthesis (or similar), and
                # plainly otherwise.
                # Can also ether strings
                match c:
                    case '(':
                        yield color(depth(1), c)
                        paren += 1
                    case ')':
                        yield color(depth(-1), c)
                        paren -= 1
                    case '[':
                        yield color(depth(1), c)
                        brack += 1
                    case ']':
                        yield color(depth(-1), c)
                        brack -= 1
                    case '{':
                        yield color(depth(1), c)
                        brace += 1
                    case '}':
                        yield color(depth(-1), c)
                        brace -= 1
                    case "'" | '"':
                        in_string = c
                        yield c
                    case '\\':
                        yield c
                        yield strm.read(1)
                    case c:
                        yield c
            case [_, _]:
                yield c