Add an ANSI color lexer

fluffy/component/ansi_color.py

@@ -0,0 +1,103 @@
+import re
+
+import pygments.lexer
+import pygments.token
+from pygments.token import Text
+
+
+# Pygments doesn't have a generic "color" token; instead everything is
+# contextual (e.g. "comment" or "variable"). That doesn't make sense for us,
+# where the actual colors actually *are* what we care about.
+Color = pygments.token.Token.Color
+
+_ansi_code_to_color = {
+    0: 'Black',
+    1: 'Red',
+    2: 'Green',
+    3: 'Yellow',
+    4: 'Blue',
+    5: 'Magenta',
+    6: 'Cyan',
+    7: 'White',
+}
+
+
+def token_from_lexer_state(bold, fg_color, bg_color):
+    """Construct a token given the current lexer state.
+
+    We can only emit one token even though we have a multiple-tuple state.
+
+    To do this, we construct tokens like "BoldRed".
+    """
+    token_name = ''
+
+    if bold:
+        token_name += 'Bold'
+
+    if fg_color:
+        token_name += fg_color
+
+    if bg_color:
+        token_name += 'BG' + bg_color
+
+    if token_name == '':
+        return Text
+    else:
+        print(token_name)
+        return getattr(Color, token_name)
+
+
+def callback(lexer, match):
+    codes = match.group(1).split(';')
+
+    for code in codes:
+        try:
+            code = int(code)
+        except ValueError:
+            # This shouldn't ever happen if we're given valid ANSI color codes,
+            # but people can pastebin junk, and we should tolerate that.
+            continue
+        else:
+            fg_color = _ansi_code_to_color.get(code - 30)
+            bg_color = _ansi_code_to_color.get(code - 40)
+            if fg_color:
+                lexer.fg_color = fg_color
+            elif bg_color:
+                lexer.bg_color = bg_color
+            elif code == 1:
+                lexer.bold = True
+            elif code == 22:
+                lexer.bold = False
+            elif code == 39:
+                lexer.fg_color = None
+            elif code == 49:
+                lexer.bg_color = None
+            elif code == 0:
+                lexer.reset_state()
+
+    yield match.start(), lexer.current_token(), match.group(2)
+
+
+class AnsiColorLexer(pygments.lexer.RegexLexer):
+    name = 'ANSI Terminal'
+    flags = re.DOTALL | re.MULTILINE
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.reset_state()
+
+    def current_token(self):
+        return token_from_lexer_state(
+            self.bold, self.fg_color, self.bg_color,
+        )
+
+    def reset_state(self):
+        self.bold = False
+        self.fg_color = None
+        self.bg_color = None
+
+    tokens = {
+        'root': [
+            (r'\x1b\[(.*?)m([^\x1b]*)', callback),
+        ],
+    }

fluffy/component/highlighting.py

@@ -1,12 +1,16 @@
+import itertools
 import re
 from collections import namedtuple
 
 import pygments
 import pygments.lexers
+import pygments.styles.xcode
 from pygments.formatters import HtmlFormatter
-from pygments.styles import get_style_by_name
+from pygments.token import Text
 from pyquery import PyQuery as pq
 
+from fluffy.component import ansi_color
+
 
 # We purposefully don't list all possible languages, and instead just the ones
 # we think people are most likely to use.
@@ -36,12 +40,56 @@
     'yaml': 'YAML',
 }
 
+FG_COLORS = {
+    'Black': '#000000',
+    'Red': '#EF2929',
+    'Green': '#62ca00',
+    'Yellow': '#dac200',
+    'Blue': '#3465A4',
+    'Magenta': '#ce42be',
+    'Cyan': '#34E2E2',
+    'White': '#ffffff',
+}
+
+BG_COLORS = {
+    'Black': '#000000',
+    'Red': '#EF2929',
+    'Green': '#8AE234',
+    'Yellow': '#FCE94F',
+    'Blue': '#3465A4',
+    'Magenta': '#c509c5',
+    'Cyan': '#34E2E2',
+    'White': '#ffffff',
+}
+
+
+class FluffyStyle(pygments.styles.xcode.XcodeStyle):
+
+    styles = dict(pygments.styles.xcode.XcodeStyle.styles)
+
+    # Add ANSI color token definitions.
+    for bold, fg_color, bg_color in itertools.product(
+        (False, True),
+        {None} | set(FG_COLORS),
+        {None} | set(BG_COLORS),
+    ):
+        token = ansi_color.token_from_lexer_state(bold, fg_color, bg_color)
+        if token is not Text:
+            value = ''
+            if bold:
+                value += ' bold'
+            if fg_color:
+                value += ' ' + FG_COLORS[fg_color]
+            if bg_color:
+                value += ' bg:' + BG_COLORS[bg_color]
+            styles[token] = value.strip()
+
 
 _pygments_formatter = HtmlFormatter(
     noclasses=True,
     linespans='line',
     nobackground=True,
-    style=get_style_by_name('xcode'),
+    style=FluffyStyle,
 )
 
 
@@ -52,7 +100,8 @@ def name(self):
         return self.lexer.name
 
     def highlight(self, text):
-        return _highlight(text, self.lexer)
+        text = _highlight(text, self.lexer)
+        return text
 
 
 class DiffHighlighter(namedtuple('PygmentsHighlighter', ('lexer',))):
@@ -94,6 +143,11 @@ def looks_like_diff(text):
     return bool(re.search(r'^diff --git ', text, re.MULTILINE))
 
 
+def looks_like_ansi_color(text):
+    """Return whether the text looks like it has ANSI color codes."""
+    return '\x1b[' in text
+
+
 def strip_diff_things(text):
     """Remove things from the text that make it look like a diff.
 
@@ -131,7 +185,13 @@ def get_highlighter(text, language):
     lexer = guess_lexer(text, language)
 
     diff_requested = (language or '').startswith('diff-')
+
     if (
+            language == 'ansi-color' or
+            (language is None and looks_like_ansi_color(text))
+    ):
+        lexer = ansi_color.AnsiColorLexer()
+    elif (
             diff_requested or
             lexer is None or
             language is None or

fluffy/run.py

@@ -36,6 +36,17 @@ def view_diff():
             raw_url='#raw',
         )
 
+    @app.route('/test/ansi-color')
+    def view_ansi_color():
+        text = (TESTING_DIR / 'files' / 'ansi-color').open().read()
+        return render_template(
+            'paste.html',
+            text=text,
+            highlighter=get_highlighter(text, None),
+            edit_url='#edit',
+            raw_url='#raw',
+        )
+
     @app.route('/test/markdown')
     def view_markdown():
         return render_template(

testing/files/ansi-color

@@ -0,0 +1,36 @@
+total 232K
+drwxr-xr-x    2 root root 4.0K Aug  8 14:53 bin
+drwxr-xr-x    3 root root 4.0K Aug  8 14:53 boot
+drwxr-xr-x   16 root root 6.4K Aug  7 17:01 dev
+drwxr-xr-x    2 root root 4.0K Jul  3 15:17 dist
+drwxr-xr-x    2 root root 4.0K Jul 14 16:25 dne
+drwxr-xr-x  130 root root  12K Aug 10 13:35 etc
+drwxr-xr-x    3 root root 4.0K Jul 14 16:26 foo
+drwxr-xr-x    3 root root 4.0K Jun 22 16:59 home
+lrwxrwxrwx    1 root root   30 Jun 30 16:40 initrd.img -> boot/initrd.img-4.4.0-1022-aws
+drwxr-xr-x    2 root root 4.0K Jul  3 15:17 itest
+drwxr-xr-x   18 root root 4.0K Jul  4 01:40 lib
+drwxr-xr-x    2 root root 4.0K Jul  5 08:40 lib64
+drwx------    2 root root  16K Jun 27  2016 lost+found
+drwxr-xr-x    2 root root 4.0K Jun 27  2016 media
+drwxr-xr-x    2 root root 4.0K Jun 27  2016 mnt
+drwxr-xr-x   23 root root 4.0K Feb  4  2017 nail
+drwxr-xr-x   11 root root 4.0K Feb  6  2017 opt
+dr-xr-xr-x 2121 root root    0 Jul  3 10:22 proc
+drwx------    5 root root 4.0K Aug 10 08:44 root
+drwxr-xr-x   35 root root 1.3K Aug 10 13:41 run
+drwxr-xr-x    2 root root  12K Aug  8 14:53 sbin
+lrwxrwxrwx    1 root root   13 Feb  4  2017 scratch -> /nail/scratch
+drwxr-xr-x    2 root root 4.0K Jun 27  2016 srv
+dr-xr-xr-x   13 root root    0 Jul  7 22:17 sys
+-rwxr-xr-x  1 ckuehl ckuehl  304 Nov 15  2015 test.py
+drwxrwxrwt 2523 root root 120K Aug 10 13:47 tmp
+drwxr-xr-x   15 root root 4.0K Aug  2 16:40 usr
+drwxr-xr-x   16 root root 4.0K Aug  3 13:25 var
+lrwxrwxrwx    1 root root   27 Jun 30 16:40 vmlinuz -> boot/vmlinuz-4.4.0-1022-aws
+total 1.4G
+
+hello world hello world hello world hello world hello world hello world hello world hello world
+hello world hello world hello world hello world hello world hello world hello world hello world
+hello world hello world hello world hello world hello world hello world hello world hello world
+hello world hello world hello world hello world hello world hello world hello world hello world