widgets/thread: generic pygments highlighting

1 files changed, 64 insertions, 24 deletions

diff --git a/alot/widgets/thread.py b/alot/widgets/thread.py
index af82cc25..6cffa909 100644
--- a/alot/widgets/thread.py
+++ b/alot/widgets/thread.py
@@ -288,31 +288,14 @@ class _TextPart(_MIMEPartWidget):
 
         self._fold_context = settings.get('thread_fold_context')
 
-        # color git patches using pygments
-        # FIXME this is quite ad-hoc
-        # make it more configurable, detect other patch formats
-        # and possibly other kinds of text
-        if (('x-mailer' in part.headers and
-             part.headers['x-mailer'][0].startswith('git-send-email')) or
-            part.content_type in ('text/x-diff', 'text/x-patch')):
-            # process as patch
-            if have_pygments:
-                try:
-                    lexer     = pygments.lexers.get_lexer_by_name('diff')
-                    formatter = pygments.formatters.get_formatter_by_name('terminal')
-                    text      = pygments.highlight(text, lexer, formatter)
-                except pygments.util.ClassNotFound:
-                    logging.debug('Could not get a lexer/formatter for diff highlighting')
-
-                self._lines     = ansi_term.parse_escapes_to_urwid(text, attr_text)
-                self._fold      = _Fold(0, 0, len(self._lines))
-                self._max_level = 0
-
-        if self._lines is None:
-            # process as a normal text mail body:
+        # try highlighting with pygments first
+        hilit = self._highlight_pygments(text, part, attr_text)
+        if hilit is None:
+            # fallback - process as a normal text mail body:
             # detect quoted blocks
-            self._lines, self._fold, self._max_level = \
-                    self._parse_quotes(text, attr_text, attrs_quote)
+            hilit = self._parse_quotes(text, attr_text, attrs_quote)
+
+        self._lines, self._fold, self._max_level = hilit
 
         self._body_placeholder = urwid.WidgetPlaceholder(urwid.Text(''))
 
@@ -423,6 +406,63 @@ class _TextPart(_MIMEPartWidget):
 
         return colored_lines, cur_fold, max_level
 
+    def _highlight_pygments(self, text, part, attr_text):
+        # color git patches using pygments
+        # FIXME this is quite ad-hoc
+        # make it more configurable, detect other patch formats
+        # and possibly other kinds of text
+        if not have_pygments:
+            return None
+
+        # get formatter for terminal
+        # TODO: support terminal256 and theming
+        try:
+            formatter = pygments.formatters.get_formatter_by_name('terminal')
+        except pygments.util.ClassNotFound:
+            logging.warning('Could not get pygments formatter for terminal')
+            return None
+
+        # try finding a lexer for this part
+        lexer = None
+
+        # first, try guessing by filename/content
+        try:
+            if part.filename is not None:
+                lexer = pygments.lexers.guess_lexer_for_filename(
+                            part.filename, text)
+            else:
+                lexer = pygments.lexers.guess_lexer(text)
+        except pygments.util.ClassNotFound:
+            logging.debug('No pygments lexer for filename: %s', part.filename)
+
+        # second, try the MIME type
+        try:
+            lexer = pygments.lexers.get_lexer_for_mimetype(part.content_type)
+        except pygments.util.ClassNotFound:
+            logging.debug('No pygments lexer for MIME tpye: %s', part.content_type)
+
+        # handle git-send-email patches, which are sent as text/plain
+        if ('x-mailer' in part.headers                               and
+            part.headers['x-mailer'][0].startswith('git-send-email') and
+            '\ndiff' in text):
+            try:
+                lexer     = pygments.lexers.get_lexer_by_name('diff')
+            except pygments.util.ClassNotFound:
+                logging.warning('Could not get a lexer/formatter for diff highlighting')
+
+        # FIXME: this is a hack to prevent getting a dummy lexer for plaintext
+        # that does nothing and prevents our quote detection code from running
+        # or the html lexer, which is not necessary since we already processed
+        # html with an external renderer
+        if lexer is None or 'text' in lexer.aliases or 'html' in lexer.aliases:
+            logging.info('No pygments lexer for MIME part: %s', part)
+            return None
+
+        text  = pygments.highlight(text, lexer, formatter)
+        lines = ansi_term.parse_escapes_to_urwid(text, attr_text)
+
+        return lines, _Fold(0, 0, len(lines)), 0
+
     @property
     def foldlevel(self):
         return self._fold_level