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