1"""Pygments related rendering
2"""
3
4
5import urwid
6import pygments
7import pygments.util
8from pygments.formatter import Formatter
9import time
10import urwid
11
12
13import lookatme.config as config
14
15
16LEXER_CACHE = {}
17STYLE_CACHE = {}
18FORMATTER_CACHE = {}
19
20
21def get_formatter(style_name):
22    style = get_style(style_name)
23
24    formatter, style_bg = FORMATTER_CACHE.get(style_name, (None, None))
25    if formatter is None:
26        style_bg = UrwidFormatter.findclosest(style.background_color.replace("#", ""))
27        formatter = UrwidFormatter(
28            style=style,
29            usebg=(style_bg is not None),
30        )
31        FORMATTER_CACHE[style_name] = (formatter, style_bg)
32    return formatter, style_bg
33
34
35def get_lexer(lang, default="text"):
36    lexer = LEXER_CACHE.get(lang, None)
37    if lexer is None:
38        try:
39            lexer = pygments.lexers.get_lexer_by_name(lang)
40        except pygments.util.ClassNotFound:
41            lexer = pygments.lexers.get_lexer_by_name(default)
42        LEXER_CACHE[lang] = lexer
43    return lexer
44
45
46def get_style(style_name):
47    style = STYLE_CACHE.get(style_name, None)
48    if style is None:
49        style = pygments.styles.get_style_by_name(style_name)
50        STYLE_CACHE[style_name] = style
51    return style
52
53
54def render_text(text, lang="text", style_name=None, plain=False):
55    """Render the provided text with the pygments renderer
56    """
57    if style_name is None:
58        style_name = config.STYLE["style"]
59
60    lexer = get_lexer(lang)
61    formatter, style_bg = get_formatter(style_name)
62
63    start = time.time()
64    code_tokens = lexer.get_tokens(text)
65    config.LOG.debug(f"Took {time.time()-start}s to render {len(text)} bytes")
66
67    markup = []
68    for x in formatter.formatgenerator(code_tokens):
69        if style_bg:
70            x[0].background = style_bg
71        markup.append(x)
72
73    if markup[-1][1] == "\n":
74        markup = markup[:-1]
75
76    if len(markup) == 0:
77        markup = [(None, "")]
78    elif markup[-1][1].endswith("\n"):
79        markup[-1] = (markup[-1][0], markup[-1][1][:-1])
80
81    if plain:
82        return markup
83    else:
84        return urwid.AttrMap(urwid.Text(markup), urwid.AttrSpec("default", style_bg))
85
86
87class UrwidFormatter(Formatter):
88    """Formatter that returns [(text,attrspec), ...],
89    where text is a piece of text, and attrspec is an urwid.AttrSpec"""
90    def __init__(self, **options):
91        """Extra arguments:
92
93        usebold: if false, bold will be ignored and always off
94                default: True
95        usebg: if false, background color will always be 'default'
96                default: True
97        colors: number of colors to use (16, 88, or 256)
98                default: 256"""
99        self.usebold = options.get('usebold',True)
100        self.usebg = options.get('usebg', True)
101        colors = options.get('colors', 256)
102        self.style_attrs = {}
103        Formatter.__init__(self, **options)
104
105    @property
106    def style(self):
107        return self._style
108
109    @style.setter
110    def style(self, newstyle):
111        self._style = newstyle
112        self._setup_styles()
113
114    @staticmethod
115    def _distance(col1, col2):
116        r1, g1, b1 = col1
117        r2, g2, b2 = col2
118
119        rd = r1 - r2
120        gd = g1 - g2
121        bd = b1 - b2
122
123        return rd*rd + gd*gd + bd*bd
124
125    @classmethod
126    def findclosest(cls, colstr, colors=256):
127        """Takes a hex string and finds the nearest color to it.
128
129        Returns a string urwid will recognize."""
130
131        rgb = int(colstr, 16)
132        r = (rgb >> 16) & 0xff
133        g = (rgb >> 8) & 0xff
134        b = rgb & 0xff
135
136        dist = 257 * 257 * 3
137        bestcol = urwid.AttrSpec('h0','default')
138
139        for i in range(colors):
140            curcol = urwid.AttrSpec('h%d' % i,'default', colors=colors)
141            currgb = curcol.get_rgb_values()[:3]
142            curdist = cls._distance((r,g,b), currgb)
143            if curdist < dist:
144                dist = curdist
145                bestcol = curcol
146
147        return bestcol.foreground
148
149    def findclosestattr(self, fgcolstr=None, bgcolstr=None, othersettings='', colors = 256):
150        """Takes two hex colstring (e.g. 'ff00dd') and returns the
151        nearest urwid style."""
152        fg = bg = 'default'
153        if fgcolstr:
154            fg = self.findclosest(fgcolstr, colors)
155        if bgcolstr:
156            bg = self.findclosest(bgcolstr, colors)
157        if othersettings:
158            fg = fg + ',' + othersettings
159        return urwid.AttrSpec(fg, bg, colors)
160
161    def _setup_styles(self, colors = 256):
162        """Fills self.style_attrs with urwid.AttrSpec attributes
163        corresponding to the closest equivalents to the given style."""
164        for ttype, ndef in self.style:
165            fgcolstr = bgcolstr = None
166            othersettings = ''
167            if ndef['color']:
168                fgcolstr = ndef['color']
169            if self.usebg and ndef['bgcolor']:
170                bgcolstr = ndef['bgcolor']
171            if self.usebold and ndef['bold']:
172                othersettings = 'bold'
173            self.style_attrs[str(ttype)] = self.findclosestattr(
174                fgcolstr, bgcolstr, othersettings, colors)
175
176    def formatgenerator(self, tokensource):
177        """Takes a token source, and generates
178        (tokenstring, urwid.AttrSpec) pairs"""
179        for (ttype, tstring) in tokensource:
180            parts = str(ttype).split(".")
181            while str(ttype) not in self.style_attrs:
182                parts = parts[:-1]
183                ttype = ".".join(parts)
184
185            attr = self.style_attrs[str(ttype)]
186            yield attr, tstring
187
188    def format(self, tokensource, outfile):
189        for (attr, tstring) in self.formatgenerator(tokensource):
190            outfile.write(attr, tstring)
191