• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2020 The Pigweed Authors
2#
3# Licensed under the Apache License, Version 2.0 (the "License"); you may not
4# use this file except in compliance with the License. You may obtain a copy of
5# the License at
6#
7#     https://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, WITHOUT
11# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
12# License for the specific language governing permissions and limitations under
13# the License.
14"""Color codes for use by rest of pw_cli."""
15
16import ctypes
17import os
18import sys
19
20import pw_cli.env
21
22
23def _make_color(*codes):
24    # Apply all the requested ANSI color codes. Note that this is unbalanced
25    # with respect to the reset, which only requires a '0' to erase all codes.
26    start = ''.join(f'\033[{code}m' for code in codes)
27    reset = '\033[0m'
28
29    return lambda msg: f'{start}{msg}{reset}'
30
31
32# TODO(keir): Replace this with something like the 'colorful' module.
33class _Color:
34    # pylint: disable=too-few-public-methods
35    # pylint: disable=too-many-instance-attributes
36    """Helpers to surround text with ASCII color escapes"""
37
38    def __init__(self):
39        self.none = str
40        self.red = _make_color(31, 1)
41        self.bold_red = _make_color(30, 41)
42        self.yellow = _make_color(33, 1)
43        self.bold_yellow = _make_color(30, 43, 1)
44        self.green = _make_color(32)
45        self.bold_green = _make_color(30, 42)
46        self.blue = _make_color(34, 1)
47        self.cyan = _make_color(36, 1)
48        self.magenta = _make_color(35, 1)
49        self.bold_magenta = _make_color(30, 45)
50        self.bold_white = _make_color(37, 1)
51        self.black_on_white = _make_color(30, 47)  # black fg white bg
52        self.black_on_green = _make_color(30, 42)  # black fg green bg
53        self.black_on_red = _make_color(30, 41)  # black fg red bg
54
55
56class _NoColor:
57    """Fake version of the _Color class that doesn't colorize."""
58
59    def __getattr__(self, _):
60        return str
61
62
63def is_enabled():
64    env = pw_cli.env.pigweed_environment()
65    # Checking if PW_USE_COLOR is in os.environ and not env since it's always
66    # in env. If it's in os.environ then use the value retrieved in env.
67    if 'PW_USE_COLOR' in os.environ:
68        return env.PW_USE_COLOR
69
70    # These are semi-standard ways to turn colors off or on for many projects.
71    # See https://bixense.com/clicolors/ and https://no-color.org/ for more.
72    if 'NO_COLOR' in os.environ:
73        return False
74    if 'CLICOLOR_FORCE' in os.environ:
75        return os.environ['CLICOLOR_FORCE'] != '0'
76
77    return sys.stdout.isatty() and sys.stderr.isatty()
78
79
80def colors(enabled: bool | None = None) -> _Color | _NoColor:
81    """Returns an object for colorizing strings.
82
83    By default, the object only colorizes if both stderr and stdout are TTYs.
84    """
85    if enabled is None:
86        enabled = is_enabled()
87
88    if enabled and os.name == 'nt':
89        # Enable ANSI color codes in Windows cmd.exe.
90        kernel32 = ctypes.windll.kernel32  # type: ignore
91        kernel32.SetConsoleMode(kernel32.GetStdHandle(-11), 7)
92
93    return _Color() if enabled else _NoColor()
94