Fix lukas-blecher#184, add highlight

Author: Freed-Wu

pix2tex/cli.py

@@ -3,6 +3,7 @@
 from PIL import ImageGrab
 from PIL import Image
 import os
+import sys
 from typing import Tuple
 import logging
 import yaml
@@ -133,7 +134,28 @@ def __call__(self, img=None, resize=True) -> str:
 
 
 def output_prediction(pred, args):
-    print(pred, '\n')
+    TERM = os.getenv('TERM', 'xterm')
+    if not sys.stdout.isatty():
+        TERM = 'dumb'
+    try:
+        from pygments import highlight
+        from pygments.lexers import get_lexer_by_name
+        from pygments.formatters import get_formatter_by_name
+
+        if TERM.split('-')[-1] == '256color':
+            formatter_name = 'terminal256'
+        elif TERM != 'dumb':
+            formatter_name = 'terminal'
+        else:
+            formatter_name = None
+        if formatter_name:
+            formatter = get_formatter_by_name(formatter_name)
+            lexer = get_lexer_by_name('tex')
+            print(highlight(pred, lexer, formatter), end='')
+    except ImportError:
+        TERM = 'dumb'
+    if TERM == 'dumb':
+        print(pred)
     if args.show or args.katex:
         try:
             if args.katex:

setup.py

@@ -24,6 +24,7 @@
     'torchtext>=0.6.0',
     'imagesize>=1.2.0',
 ]
+highlight = ['pygments']
 
 setuptools.setup(
     name='pix2tex',
@@ -66,10 +67,11 @@
         'albumentations>=0.5.2',
     ],
     extras_require={
-        'all': gui+api+train,
+        'all': gui+api+train+highlight,
         'gui': gui,
         'api': api,
-        'train': train
+        'train': train,
+        'highlight': highlight,
     },
     entry_points={
         'console_scripts': [