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"""Defines ANSI color codes."""
15
16import ctypes
17import os
18
19
20def _make_color(*codes):
21    # Apply all the requested ANSI color codes. Note that this is unbalanced
22    # with respect to the reset, which only requires a '0' to erase all codes.
23    start = ''.join('\033[{}m'.format(code) for code in codes)
24    reset = '\033[0m'
25
26    return staticmethod(lambda msg: u'{}{}{}'.format(start, msg, reset))
27
28
29class Color:  # pylint: disable=too-few-public-methods
30    """Helpers to surround text with ASCII color escapes"""
31    bold = _make_color(1)
32    red = _make_color(31)
33    bold_red = _make_color(31, 1)
34    green = _make_color(32)
35    magenta = _make_color(35, 1)
36
37
38def enable_colors():
39    if os.name == 'nt':
40        kernel32 = ctypes.windll.kernel32
41        kernel32.SetConsoleMode(kernel32.GetStdHandle(-11), 7)
42