1# Copyright 2021 gRPC authors.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14"""The module contains helpers to enable color output in terminals.
15
16Use this to log resources dumped as a structured document (f.e. YAML),
17and enable colorful syntax highlighting.
18
19TODO(sergiitk): This can be used to output protobuf responses formatted as JSON.
20"""
21import logging
22from typing import Optional
23
24from absl import flags
25import pygments
26import pygments.formatter
27import pygments.formatters.other
28import pygments.formatters.terminal
29import pygments.formatters.terminal256
30import pygments.lexer
31import pygments.lexers.data
32import pygments.styles
33
34# The style for terminals supporting 8/16 colors.
35STYLE_ANSI_16 = 'ansi16'
36# Join with pygments styles for terminals supporting 88/256 colors.
37ALL_COLOR_STYLES = [STYLE_ANSI_16] + list(pygments.styles.get_all_styles())
38
39# Flags.
40COLOR = flags.DEFINE_bool("color", default=True, help='Colorize the output')
41COLOR_STYLE = flags.DEFINE_enum(
42    "color_style",
43    default='material',
44    enum_values=ALL_COLOR_STYLES,
45    help=('Color styles for terminals supporting 256 colors. '
46          f'Use {STYLE_ANSI_16} style for terminals supporting 8/16 colors'))
47
48logger = logging.getLogger(__name__)
49
50# Type aliases.
51Lexer = pygments.lexer.Lexer
52YamlLexer = pygments.lexers.data.YamlLexer
53Formatter = pygments.formatter.Formatter
54NullFormatter = pygments.formatters.other.NullFormatter
55TerminalFormatter = pygments.formatters.terminal.TerminalFormatter
56Terminal256Formatter = pygments.formatters.terminal256.Terminal256Formatter
57
58
59class Highlighter:
60    formatter: Formatter
61    lexer: Lexer
62    color: bool
63    color_style: Optional[str] = None
64
65    def __init__(self,
66                 *,
67                 lexer: Lexer,
68                 color: Optional[bool] = None,
69                 color_style: Optional[str] = None):
70        self.lexer = lexer
71        self.color = color if color is not None else COLOR.value
72
73        if self.color:
74            color_style = color_style if color_style else COLOR_STYLE.value
75            if color_style not in ALL_COLOR_STYLES:
76                raise ValueError(f'Unrecognized color style {color_style}, '
77                                 f'valid styles: {ALL_COLOR_STYLES}')
78            if color_style == STYLE_ANSI_16:
79                # 8/16 colors support only.
80                self.formatter = TerminalFormatter()
81            else:
82                # 88/256 colors.
83                self.formatter = Terminal256Formatter(style=color_style)
84        else:
85            self.formatter = NullFormatter()
86
87    def highlight(self, code: str) -> str:
88        return pygments.highlight(code, self.lexer, self.formatter)
89
90
91class HighlighterYaml(Highlighter):
92
93    def __init__(self,
94                 *,
95                 color: Optional[bool] = None,
96                 color_style: Optional[str] = None):
97        super().__init__(lexer=YamlLexer(encoding='utf-8'),
98                         color=color,
99                         color_style=color_style)
100