1# --------------------------------------------------------------------------------------------
2# Copyright (c) Microsoft Corporation. All rights reserved.
3# Licensed under the MIT License. See License.txt in the project root for license information.
4# --------------------------------------------------------------------------------------------
5
6"""
7Support styled output.
8
9Currently, only color is supported, underline/bold/italic may be supported in the future.
10
11Design spec:
12https://devdivdesignguide.azurewebsites.net/command-line-interface/color-guidelines-for-command-line-interface/
13
14Console Virtual Terminal Sequences:
15https://docs.microsoft.com/en-us/windows/console/console-virtual-terminal-sequences#text-formatting
16
17For a complete demo, see `src/azure-cli/azure/cli/command_modules/util/custom.py` and run `az demo style`.
18"""
19
20import sys
21from enum import Enum
22
23from knack.log import get_logger
24from knack.util import is_modern_terminal
25
26
27logger = get_logger(__name__)
28
29
30class Style(str, Enum):
31    PRIMARY = "primary"
32    SECONDARY = "secondary"
33    IMPORTANT = "important"
34    ACTION = "action"  # name TBD
35    HYPERLINK = "hyperlink"
36    # Message colors
37    ERROR = "error"
38    SUCCESS = "success"
39    WARNING = "warning"
40
41
42def _rgb_hex(rgb_hex: str):
43    """
44    Convert RGB hex value to Control Sequences.
45    """
46    template = '\x1b[38;2;{r};{g};{b}m'
47    if rgb_hex.startswith("#"):
48        rgb_hex = rgb_hex[1:]
49
50    rgb = {}
51    for i, c in enumerate(('r', 'g', 'b')):
52        value_str = rgb_hex[i * 2: i * 2 + 2]
53        value_int = int(value_str, 16)
54        rgb[c] = value_int
55
56    return template.format(**rgb)
57
58
59DEFAULT = '\x1b[0m'  # Default
60
61# Theme that doesn't contain any style
62THEME_NONE = None
63
64# Theme to be used in a dark-themed terminal
65THEME_DARK = {
66    Style.PRIMARY: DEFAULT,
67    Style.SECONDARY: '\x1b[90m',  # Bright Foreground Black
68    Style.IMPORTANT: '\x1b[95m',  # Bright Foreground Magenta
69    Style.ACTION: '\x1b[94m',  # Bright Foreground Blue
70    Style.HYPERLINK: '\x1b[96m',  # Bright Foreground Cyan
71    Style.ERROR: '\x1b[91m',  # Bright Foreground Red
72    Style.SUCCESS: '\x1b[92m',  # Bright Foreground Green
73    Style.WARNING: '\x1b[93m'  # Bright Foreground Yellow
74}
75
76# Theme to be used in a light-themed terminal
77THEME_LIGHT = {
78    Style.PRIMARY: DEFAULT,
79    Style.SECONDARY: '\x1b[90m',  # Bright Foreground Black
80    Style.IMPORTANT: '\x1b[35m',  # Foreground Magenta
81    Style.ACTION: '\x1b[34m',  # Foreground Blue
82    Style.HYPERLINK: '\x1b[36m',  # Foreground Cyan
83    Style.ERROR: '\x1b[31m',  # Foreground Red
84    Style.SUCCESS: '\x1b[32m',  # Foreground Green
85    Style.WARNING: '\x1b[33m'  # Foreground Yellow
86}
87
88# Theme to be used in Cloud Shell
89# Text and background's Contrast Ratio should be above 4.5:1
90THEME_CLOUD_SHELL = {
91    Style.PRIMARY: _rgb_hex('#ffffff'),
92    Style.SECONDARY: _rgb_hex('#bcbcbc'),
93    Style.IMPORTANT: _rgb_hex('#f887ff'),
94    Style.ACTION: _rgb_hex('#6cb0ff'),
95    Style.HYPERLINK: _rgb_hex('#72d7d8'),
96    Style.ERROR: _rgb_hex('#f55d5c'),
97    Style.SUCCESS: _rgb_hex('#70d784'),
98    Style.WARNING: _rgb_hex('#fbd682'),
99}
100
101
102class Theme(str, Enum):
103    DARK = 'dark'
104    LIGHT = 'light'
105    CLOUD_SHELL = 'cloud-shell'
106    NONE = 'none'
107
108
109THEME_DEFINITIONS = {
110    Theme.DARK: THEME_DARK,
111    Theme.LIGHT: THEME_LIGHT,
112    Theme.CLOUD_SHELL: THEME_CLOUD_SHELL,
113    Theme.NONE: THEME_NONE
114}
115
116# Blue and bright blue is not visible under the default theme of powershell.exe
117POWERSHELL_COLOR_REPLACEMENT = {
118    '\x1b[34m': DEFAULT,  # Foreground Blue
119    '\x1b[94m': DEFAULT  # Bright Foreground Blue
120}
121
122
123def print_styled_text(*styled_text_objects, file=None, **kwargs):
124    """
125    Print styled text. This function wraps the built-in function `print`, additional arguments can be sent
126    via keyword arguments.
127
128    :param styled_text_objects: The input text objects. See format_styled_text for formats of each object.
129    :param file: The file to print the styled text. The default target is sys.stderr.
130    """
131    formatted_list = [format_styled_text(obj) for obj in styled_text_objects]
132    # Always fetch the latest sys.stderr in case it has been wrapped by colorama.
133    print(*formatted_list, file=file or sys.stderr, **kwargs)
134
135
136def format_styled_text(styled_text, theme=None):
137    """Format styled text. Dark theme used by default. Available themes are 'dark', 'light', 'none'.
138
139    To change theme for all invocations of this function, set `format_styled_text.theme`.
140    To change theme for one invocation, set parameter `theme`.
141
142    :param styled_text: Can be in these formats:
143        - text
144        - (style, text)
145        - [(style, text), ...]
146    :param theme: The theme used to format text. Can be theme name str, `Theme` Enum or dict.
147    """
148    if theme is None:
149        theme = getattr(format_styled_text, "theme", THEME_DARK)
150
151    # Convert str to the theme dict
152    if isinstance(theme, str):
153        theme = get_theme_dict(theme)
154
155    # If style is enabled, cache the value of is_legacy_powershell.
156    # Otherwise if theme is None, is_legacy_powershell is meaningless.
157    is_legacy_powershell = None
158    if theme:
159        if not hasattr(format_styled_text, "_is_legacy_powershell"):
160            from azure.cli.core.util import get_parent_proc_name
161            is_legacy_powershell = not is_modern_terminal() and get_parent_proc_name() == "powershell.exe"
162            setattr(format_styled_text, "_is_legacy_powershell", is_legacy_powershell)
163        is_legacy_powershell = getattr(format_styled_text, "_is_legacy_powershell")
164
165    # https://python-prompt-toolkit.readthedocs.io/en/stable/pages/printing_text.html#style-text-tuples
166    formatted_parts = []
167
168    # A str as PRIMARY text
169    if isinstance(styled_text, str):
170        styled_text = [(Style.PRIMARY, styled_text)]
171
172    # A tuple
173    if isinstance(styled_text, tuple):
174        styled_text = [styled_text]
175
176    for text in styled_text:
177        # str can also be indexed, bypassing IndexError, so explicitly check if the type is tuple
178        if not (isinstance(text, tuple) and len(text) == 2):
179            from azure.cli.core.azclierror import CLIInternalError
180            raise CLIInternalError("Invalid styled text. It should be a list of 2-element tuples.")
181
182        style, raw_text = text
183
184        if theme:
185            try:
186                escape_seq = theme[style]
187            except KeyError:
188                from azure.cli.core.azclierror import CLIInternalError
189                raise CLIInternalError("Invalid style. Only use pre-defined style in Style enum.")
190            # Replace blue in powershell.exe
191            if is_legacy_powershell and escape_seq in POWERSHELL_COLOR_REPLACEMENT:
192                escape_seq = POWERSHELL_COLOR_REPLACEMENT[escape_seq]
193            formatted_parts.append(escape_seq + raw_text)
194        else:
195            formatted_parts.append(raw_text)
196
197    # Reset control sequence
198    if theme is not THEME_NONE:
199        formatted_parts.append(DEFAULT)
200    return ''.join(formatted_parts)
201
202
203def highlight_command(raw_command):
204    """Highlight a command with colors.
205
206    For example, for
207
208        az group create --name myrg --location westus
209
210    The command name 'az group create', argument name '--name', '--location' are marked as ACTION style.
211    The argument value 'myrg' and 'westus' are marked as PRIMARY style.
212    If the argument is provided as '--location=westus', it will be marked as PRIMARY style.
213
214    :param raw_command: The command that needs to be highlighted.
215    :type raw_command: str
216    :return: The styled command text.
217    :rtype: list
218    """
219
220    styled_command = []
221    argument_begins = False
222
223    for index, arg in enumerate(raw_command.split()):
224        spaced_arg = ' {}'.format(arg) if index > 0 else arg
225        style = Style.PRIMARY
226
227        if arg.startswith('-') and '=' not in arg:
228            style = Style.ACTION
229            argument_begins = True
230        elif not argument_begins and '=' not in arg:
231            style = Style.ACTION
232
233        styled_command.append((style, spaced_arg))
234
235    return styled_command
236
237
238def get_theme_dict(theme: str):
239    try:
240        return THEME_DEFINITIONS[theme]
241    except KeyError as ex:
242        available_themes = ', '.join([m.value for m in Theme.__members__.values()])  # pylint: disable=no-member
243        logger.warning("Invalid theme %s. Supported themes: %s", ex, available_themes)
244        return None
245