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
|
#!/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()
|
