Skip to content

Commit 36c70b9

Browse files
committed
use an actual ANSI code formatter
1 parent 7c06f5d commit 36c70b9

2 files changed

Lines changed: 86 additions & 40 deletions

File tree

refinery/units/sinks/hl.py

Lines changed: 69 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
from __future__ import annotations
22

33
import pathlib
4+
import typing
45

56
import colorama
67

@@ -12,6 +13,10 @@
1213
from refinery.lib.types import Param
1314
from refinery.units import Arg, Unit
1415

16+
if typing.TYPE_CHECKING:
17+
from io import TextIOBase
18+
from pygments.token import _TokenType as TokenType
19+
1520

1621
class hl(Unit):
1722
"""
@@ -61,9 +66,9 @@ def __init__(
6166
@Unit.Requires('Pygments', 3)
6267
def _pygments():
6368
import pygments
69+
import pygments.formatter
6470
import pygments.formatters
6571
import pygments.lexers
66-
import pygments.style
6772
import pygments.token
6873
return pygments
6974

@@ -72,7 +77,6 @@ def _style_variant(self):
7277

7378
def process(self, data):
7479
lib = self._pygments
75-
token = lib.token
7680
language: str = self.args.language
7781

7882
if language:
@@ -109,28 +113,70 @@ def process(self, data):
109113
tf = lib.formatters.TerminalFormatter(bg='light')
110114
elif self.args.dark:
111115
tf = lib.formatters.TerminalFormatter(bg='dark')
116+
elif (style := self.args.style):
117+
tf = t256(colorscheme=style)
112118
else:
113-
if not (style := self.args.style):
114-
class _style(lib.style.Style):
115-
background_color = 'default'
116-
R = 'ansibrightred'
117-
C = 'ansibrightcyan'
118-
W = 'ansiwhite'
119-
B = 'ansibrightblack'
120-
styles = {
121-
token.Comment : B,
122-
token.Text : W,
123-
token.Name : W,
124-
token.Error : R,
125-
token.Keyword : R,
126-
token.Operator : R,
127-
token.Punctuation : R,
128-
token.String : C,
129-
token.Number : C,
130-
token.Literal : C,
131-
}
132-
style = _style
133-
tf = t256(style=style)
119+
FG = colorama.Fore
120+
ST = colorama.Style
121+
t = lib.token
122+
123+
R = FG.LIGHTRED_EX
124+
C = FG.CYAN
125+
W = FG.WHITE
126+
B = FG.LIGHTBLACK_EX
127+
128+
ANSI_MAP = {
129+
# Base layer
130+
t.Text : W,
131+
# Comments (background noise)
132+
t.Comment : B,
133+
t.Comment.Preproc : B,
134+
# Structure (dominant visual signal)
135+
t.Keyword : R + ST.BRIGHT,
136+
t.Keyword.Constant : R,
137+
t.Keyword.Type : R,
138+
t.Operator : R,
139+
t.Operator.Word : R,
140+
# Names (keep calm and readable)
141+
t.Name : W,
142+
t.Name.Variable : W,
143+
t.Name.Function : W,
144+
# Rare highlight (just enough contrast)
145+
t.Name.Class : W + ST.BRIGHT,
146+
# Keep palette tight (no extra hues)
147+
t.Name.Namespace : B,
148+
t.Name.Attribute : B,
149+
t.Name.Property : B,
150+
# Structural accents
151+
t.Name.Builtin : R,
152+
t.Name.Decorator : R,
153+
t.Name.Tag : R,
154+
# Data (controlled green usage)
155+
t.String : C,
156+
t.String.Doc : C,
157+
t.String.Escape : C,
158+
t.String.Interpol : C,
159+
t.Number : C,
160+
t.Literal : C,
161+
t.Name.Constant : C,
162+
# Background structure
163+
t.Punctuation : B,
164+
# Errors (high visibility)
165+
t.Error : R + ST.BRIGHT,
166+
}
167+
168+
class ColoramaFormatter(lib.formatter.Formatter):
169+
def format(self, tokensource: list[tuple[TokenType, str]], outfile: TextIOBase):
170+
for tt, value in tokensource:
171+
cc = None
172+
while tt and cc is None:
173+
cc = ANSI_MAP.get(tt)
174+
tt = tt.parent
175+
outfile.write(cc or FG.WHITE)
176+
outfile.write(value)
177+
outfile.write(ST.RESET_ALL)
178+
179+
tf = ColoramaFormatter()
134180

135181
out = lib.highlight(data, lexer, tf)
136182
out = F'{out}{colorama.Style.RESET_ALL}'

test/units/sinks/test_hl.py

Lines changed: 17 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -16,35 +16,35 @@ def test_gruvbox_python(self):
1616
unit = self.load('py')
1717
result = data | unit | bytes
1818
self.assertEqual(result, bytes.fromhex(
19-
'1B5B 3031 6D70 7269 6E74 1B5B 3339 6D1B 5B39 316D 281B 5B33 396D 1B5B 3936 6D22'
20-
'1B5B 3339 6D1B 5B39 366D 6865 6C6C 6F1B 5B33 396D 1B5B 3936 6D22 1B5B 3339 6D1B'
21-
'5B39 316D 291B 5B33 396D 0A1B 5B30 6D'
19+
'1B5B 3931 6D70 7269 6E74 1B5B 306D 1B5B 3930 6D28 1B5B 306D 1B5B 3336 6D22 1B5B'
20+
'306D 1B5B 3336 6D68 656C 6C6F 1B5B 306D 1B5B 3336 6D22 1B5B 306D 1B5B 3930 6D29'
21+
'1B5B 306D 1B5B 3337 6D0A 1B5B 306D 1B5B 306D'
2222
))
2323

2424
def test_github_html(self):
2525
data = b'<html><body>Hello</body></html>'
2626
unit = self.load('html')
2727
result = data | unit | bytes
2828
self.assertEqual(result, bytes.fromhex(
29-
'1B5B 3931 6D3C 1B5B 3339 6D1B 5B30 316D 6874 6D6C 1B5B 3339 6D1B 5B39 316D 3E1B'
30-
'5B33 396D 1B5B 3931 6D3C 1B5B 3339 6D1B 5B30 316D 626F 6479 1B5B 3339 6D1B 5B39'
31-
'316D 3E1B 5B33 396D 1B5B 3031 6D48 656C 6C6F 1B5B 3339 6D1B 5B39 316D 3C1B 5B33'
32-
'396D 1B5B 3931 6D2F 1B5B 3339 6D1B 5B30 316D 626F 6479 1B5B 3339 6D1B 5B39 316D'
33-
'3E1B 5B33 396D 1B5B 3931 6D3C 1B5B 3339 6D1B 5B39 316D 2F1B 5B33 396D 1B5B 3031'
34-
'6D68 746D 6C1B 5B33 396D 1B5B 3931 6D3E 1B5B 3339 6D0A 1B5B 306D'
29+
'1B5B 3930 6D3C 1B5B 306D 1B5B 3931 6D68 746D 6C1B 5B30 6D1B 5B39 306D 3E1B 5B30'
30+
'6D1B 5B39 306D 3C1B 5B30 6D1B 5B39 316D 626F 6479 1B5B 306D 1B5B 3930 6D3E 1B5B'
31+
'306D 1B5B 3337 6D48 656C 6C6F 1B5B 306D 1B5B 3930 6D3C 1B5B 306D 1B5B 3930 6D2F'
32+
'1B5B 306D 1B5B 3931 6D62 6F64 791B 5B30 6D1B 5B39 306D 3E1B 5B30 6D1B 5B39 306D'
33+
'3C1B 5B30 6D1B 5B39 306D 2F1B 5B30 6D1B 5B39 316D 6874 6D6C 1B5B 306D 1B5B 3930'
34+
'6D3E 1B5B 306D 1B5B 3337 6D0A 1B5B 306D 1B5B 306D'
3535
))
3636

3737
def test_solarized_javascript(self):
3838
data = b'if (var > 9) { console.log("Large!"); }'
3939
unit = self.load('js')
4040
result = data | unit | bytes
4141
self.assertEqual(result, bytes.fromhex(
42-
'1B5B 3931 6D69 661B 5B33 396D 1B5B 3031 6D20 1B5B 3339 6D1B 5B39 316D 281B 5B33'
43-
'396D 1B5B 3931 6D76 6172 1B5B 3339 6D1B 5B30 316D 201B 5B33 396D 1B5B 3931 6D3E'
44-
'1B5B 3339 6D1B 5B30 316D 201B 5B33 396D 1B5B 3936 6D39 1B5B 3339 6D1B 5B39 316D'
45-
'291B 5B33 396D 1B5B 3031 6D20 1B5B 3339 6D1B 5B39 316D 7B1B 5B33 396D 1B5B 3031'
46-
'6D20 1B5B 3339 6D1B 5B30 316D 636F 6E73 6F6C 651B 5B33 396D 1B5B 3931 6D2E 1B5B'
47-
'3339 6D1B 5B30 316D 6C6F 671B 5B33 396D 1B5B 3931 6D28 1B5B 3339 6D1B 5B39 366D'
48-
'224C 6172 6765 2122 1B5B 3339 6D1B 5B39 316D 291B 5B33 396D 1B5B 3931 6D3B 1B5B'
49-
'3339 6D1B 5B30 316D 201B 5B33 396D 1B5B 3931 6D7D 1B5B 3339 6D0A 1B5B 306D'
42+
'1B5B 3931 6D1B 5B31 6D69 661B 5B30 6D1B 5B33 376D 201B 5B30 6D1B 5B39 306D 281B'
43+
'5B30 6D1B 5B39 316D 1B5B 316D 7661 721B 5B30 6D1B 5B33 376D 201B 5B30 6D1B 5B39'
44+
'316D 3E1B 5B30 6D1B 5B33 376D 201B 5B30 6D1B 5B33 366D 391B 5B30 6D1B 5B39 306D'
45+
'291B 5B30 6D1B 5B33 376D 201B 5B30 6D1B 5B39 306D 7B1B 5B30 6D1B 5B33 376D 201B'
46+
'5B30 6D1B 5B33 376D 636F 6E73 6F6C 651B 5B30 6D1B 5B39 306D 2E1B 5B30 6D1B 5B33'
47+
'376D 6C6F 671B 5B30 6D1B 5B39 306D 281B 5B30 6D1B 5B33 366D 224C 6172 6765 2122'
48+
'1B5B 306D 1B5B 3930 6D29 1B5B 306D 1B5B 3930 6D3B 1B5B 306D 1B5B 3337 6D20 1B5B'
49+
'306D 1B5B 3930 6D7D 1B5B 306D 1B5B 3337 6D0A 1B5B 306D 1B5B 306D'
5050
))

0 commit comments

Comments
 (0)