1import json
2from typing import Optional, Type
3
4import pygments.lexer
5import pygments.lexers
6import pygments.style
7import pygments.styles
8import pygments.token
9from pygments.formatters.terminal import TerminalFormatter
10from pygments.formatters.terminal256 import Terminal256Formatter
11from pygments.lexer import Lexer
12from pygments.lexers.data import JsonLexer
13from pygments.lexers.special import TextLexer
14from pygments.lexers.text import HttpLexer as PygmentsHttpLexer
15from pygments.util import ClassNotFound
16
17from ..lexers.json import EnhancedJsonLexer
18from ...compat import is_windows
19from ...context import Environment
20from ...plugins import FormatterPlugin
21
22
23AUTO_STYLE = 'auto'  # Follows terminal ANSI color styles
24DEFAULT_STYLE = AUTO_STYLE
25SOLARIZED_STYLE = 'solarized'  # Bundled here
26if is_windows:
27    # Colors on Windows via colorama don't look that
28    # great and fruity seems to give the best result there.
29    DEFAULT_STYLE = 'fruity'
30
31AVAILABLE_STYLES = set(pygments.styles.get_all_styles())
32AVAILABLE_STYLES.add(SOLARIZED_STYLE)
33AVAILABLE_STYLES.add(AUTO_STYLE)
34
35
36class ColorFormatter(FormatterPlugin):
37    """
38    Colorize using Pygments
39
40    This processor that applies syntax highlighting to the headers,
41    and also to the body if its content type is recognized.
42
43    """
44    group_name = 'colors'
45
46    def __init__(
47        self,
48        env: Environment,
49        explicit_json=False,
50        color_scheme=DEFAULT_STYLE,
51        **kwargs
52    ):
53        super().__init__(**kwargs)
54
55        if not env.colors:
56            self.enabled = False
57            return
58
59        use_auto_style = color_scheme == AUTO_STYLE
60        has_256_colors = env.colors == 256
61        if use_auto_style or not has_256_colors:
62            http_lexer = PygmentsHttpLexer()
63            formatter = TerminalFormatter()
64        else:
65            from ..lexers.http import SimplifiedHTTPLexer
66            http_lexer = SimplifiedHTTPLexer()
67            formatter = Terminal256Formatter(
68                style=self.get_style_class(color_scheme)
69            )
70
71        self.explicit_json = explicit_json  # --json
72        self.formatter = formatter
73        self.http_lexer = http_lexer
74
75    def format_headers(self, headers: str) -> str:
76        return pygments.highlight(
77            code=headers,
78            lexer=self.http_lexer,
79            formatter=self.formatter,
80        ).strip()
81
82    def format_body(self, body: str, mime: str) -> str:
83        lexer = self.get_lexer_for_body(mime, body)
84        if lexer:
85            body = pygments.highlight(
86                code=body,
87                lexer=lexer,
88                formatter=self.formatter,
89            )
90        return body
91
92    def get_lexer_for_body(
93        self, mime: str,
94        body: str
95    ) -> Optional[Type[Lexer]]:
96        return get_lexer(
97            mime=mime,
98            explicit_json=self.explicit_json,
99            body=body,
100        )
101
102    @staticmethod
103    def get_style_class(color_scheme: str) -> Type[pygments.style.Style]:
104        try:
105            return pygments.styles.get_style_by_name(color_scheme)
106        except ClassNotFound:
107            return Solarized256Style
108
109
110def get_lexer(
111    mime: str,
112    explicit_json=False,
113    body=''
114) -> Optional[Type[Lexer]]:
115    # Build candidate mime type and lexer names.
116    mime_types, lexer_names = [mime], []
117    type_, subtype = mime.split('/', 1)
118    if '+' not in subtype:
119        lexer_names.append(subtype)
120    else:
121        subtype_name, subtype_suffix = subtype.split('+', 1)
122        lexer_names.extend([subtype_name, subtype_suffix])
123        mime_types.extend([
124            f'{type_}/{subtype_name}',
125            f'{type_}/{subtype_suffix}',
126        ])
127
128    # As a last resort, if no lexer feels responsible, and
129    # the subtype contains 'json', take the JSON lexer
130    if 'json' in subtype:
131        lexer_names.append('json')
132
133    # Try to resolve the right lexer.
134    lexer = None
135    for mime_type in mime_types:
136        try:
137            lexer = pygments.lexers.get_lexer_for_mimetype(mime_type)
138            break
139        except ClassNotFound:
140            pass
141    else:
142        for name in lexer_names:
143            try:
144                lexer = pygments.lexers.get_lexer_by_name(name)
145            except ClassNotFound:
146                pass
147
148    if explicit_json and body and (not lexer or isinstance(lexer, TextLexer)):
149        # JSON response with an incorrect Content-Type?
150        try:
151            json.loads(body)  # FIXME: the body also gets parsed in json.py
152        except ValueError:
153            pass  # Nope
154        else:
155            lexer = pygments.lexers.get_lexer_by_name('json')
156
157    # Use our own JSON lexer: it supports JSON bodies preceded by non-JSON data
158    # as well as legit JSON bodies.
159    if isinstance(lexer, JsonLexer):
160        lexer = EnhancedJsonLexer()
161
162    return lexer
163
164
165class Solarized256Style(pygments.style.Style):
166    """
167    solarized256
168    ------------
169
170    A Pygments style inspired by Solarized's 256 color mode.
171
172    :copyright: (c) 2011 by Hank Gay, (c) 2012 by John Mastro.
173    :license: BSD, see LICENSE for more details.
174
175    """
176    BASE03 = "#1c1c1c"
177    BASE02 = "#262626"
178    BASE01 = "#4e4e4e"
179    BASE00 = "#585858"
180    BASE0 = "#808080"
181    BASE1 = "#8a8a8a"
182    BASE2 = "#d7d7af"
183    BASE3 = "#ffffd7"
184    YELLOW = "#af8700"
185    ORANGE = "#d75f00"
186    RED = "#af0000"
187    MAGENTA = "#af005f"
188    VIOLET = "#5f5faf"
189    BLUE = "#0087ff"
190    CYAN = "#00afaf"
191    GREEN = "#5f8700"
192
193    background_color = BASE03
194    styles = {
195        pygments.token.Keyword: GREEN,
196        pygments.token.Keyword.Constant: ORANGE,
197        pygments.token.Keyword.Declaration: BLUE,
198        pygments.token.Keyword.Namespace: ORANGE,
199        pygments.token.Keyword.Reserved: BLUE,
200        pygments.token.Keyword.Type: RED,
201        pygments.token.Name.Attribute: BASE1,
202        pygments.token.Name.Builtin: BLUE,
203        pygments.token.Name.Builtin.Pseudo: BLUE,
204        pygments.token.Name.Class: BLUE,
205        pygments.token.Name.Constant: ORANGE,
206        pygments.token.Name.Decorator: BLUE,
207        pygments.token.Name.Entity: ORANGE,
208        pygments.token.Name.Exception: YELLOW,
209        pygments.token.Name.Function: BLUE,
210        pygments.token.Name.Tag: BLUE,
211        pygments.token.Name.Variable: BLUE,
212        pygments.token.String: CYAN,
213        pygments.token.String.Backtick: BASE01,
214        pygments.token.String.Char: CYAN,
215        pygments.token.String.Doc: CYAN,
216        pygments.token.String.Escape: RED,
217        pygments.token.String.Heredoc: CYAN,
218        pygments.token.String.Regex: RED,
219        pygments.token.Number: CYAN,
220        pygments.token.Operator: BASE1,
221        pygments.token.Operator.Word: GREEN,
222        pygments.token.Comment: BASE01,
223        pygments.token.Comment.Preproc: GREEN,
224        pygments.token.Comment.Special: GREEN,
225        pygments.token.Generic.Deleted: CYAN,
226        pygments.token.Generic.Emph: 'italic',
227        pygments.token.Generic.Error: RED,
228        pygments.token.Generic.Heading: ORANGE,
229        pygments.token.Generic.Inserted: GREEN,
230        pygments.token.Generic.Strong: 'bold',
231        pygments.token.Generic.Subheading: ORANGE,
232        pygments.token.Token: BASE1,
233        pygments.token.Token.Other: ORANGE,
234    }
235