-
Notifications
You must be signed in to change notification settings - Fork 6
/
setup_color.py
133 lines (112 loc) · 5.05 KB
/
setup_color.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
#!/usr/bin/env python3
# Copyright (c) 2016 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
import ctypes
import platform
import subprocess
import sys
from third_party import colorama
IS_TTY = None
OUT_TYPE = 'unknown'
def enable_native_ansi():
"""Enables native ANSI sequences in console. Windows 10 only.
Returns whether successful.
"""
kernel32 = ctypes.windll.kernel32
ENABLE_VIRTUAL_TERMINAL_PROCESSING = 0x04
out_handle = kernel32.GetStdHandle(subprocess.STD_OUTPUT_HANDLE)
# GetConsoleMode fails if the terminal isn't native.
mode = ctypes.wintypes.DWORD()
if kernel32.GetConsoleMode(out_handle, ctypes.byref(mode)) == 0:
return False
if not (mode.value & ENABLE_VIRTUAL_TERMINAL_PROCESSING):
if kernel32.SetConsoleMode(
out_handle,
mode.value | ENABLE_VIRTUAL_TERMINAL_PROCESSING) == 0:
print('kernel32.SetConsoleMode to enable ANSI sequences failed',
file=sys.stderr)
return False
return True
def init():
# should_wrap instructs colorama to wrap stdout/stderr with an ANSI
# colorcode interpreter that converts them to SetConsoleTextAttribute calls.
# This only should be True in cases where we're connected to cmd.exe's
# console. Setting this to True on non-windows systems has no effect.
should_wrap = False
global IS_TTY, OUT_TYPE
IS_TTY = sys.stdout.isatty()
is_windows = sys.platform.startswith('win')
if IS_TTY:
# Yay! We detected a console in the normal way. It doesn't really matter
# if it's windows or not, we win.
OUT_TYPE = 'console'
should_wrap = True
elif is_windows:
# assume this is some sort of file
OUT_TYPE = 'file (win)'
import msvcrt
h = msvcrt.get_osfhandle(sys.stdout.fileno())
# h is the win32 HANDLE for stdout.
ftype = ctypes.windll.kernel32.GetFileType(h)
if ftype == 2: # FILE_TYPE_CHAR
# This is a normal cmd console, but we'll only get here if we're
# running inside a `git command` which is actually
# git->bash->command. Not sure why isatty doesn't detect this case.
OUT_TYPE = 'console (cmd via msys)'
IS_TTY = True
should_wrap = True
elif ftype == 3: # FILE_TYPE_PIPE
OUT_TYPE = 'pipe (win)'
# This is some kind of pipe on windows. This could either be a real
# pipe or this could be msys using a pipe to emulate a pty. We use
# the same algorithm that msys-git uses to determine if it's
# connected to a pty or not.
# This function and the structures are defined in the MSDN
# documentation using the same names.
def NT_SUCCESS(status):
# The first two bits of status are the severity. The success
# severities are 0 and 1, and the !success severities are 2 and
# 3. Therefore since ctypes interprets the default restype of
# the call to be an 'C int' (which is guaranteed to be signed 32
# bits), All success codes are positive, and all !success codes
# are negative.
return status >= 0
class UNICODE_STRING(ctypes.Structure):
_fields_ = [('Length', ctypes.c_ushort),
('MaximumLength', ctypes.c_ushort),
('Buffer', ctypes.c_wchar_p)]
class OBJECT_NAME_INFORMATION(ctypes.Structure):
_fields_ = [('Name', UNICODE_STRING),
('NameBuffer', ctypes.c_wchar_p)]
buf = ctypes.create_string_buffer(1024)
# Ask NT what the name of the object our stdout HANDLE is. It would
# be possible to use GetFileInformationByHandleEx, but it's only
# available on Vista+. If you're reading this in 2017 or later, feel
# free to refactor this out.
#
# The '1' here is ObjectNameInformation
if NT_SUCCESS(
ctypes.windll.ntdll.NtQueryObject(h, 1, buf,
len(buf) - 2, None)):
out = OBJECT_NAME_INFORMATION.from_buffer(buf)
name = out.Name.Buffer.split('\\')[-1]
IS_TTY = name.startswith('msys-') and '-pty' in name
if IS_TTY:
OUT_TYPE = 'bash (msys)'
else:
# A normal file, or an unknown file type.
pass
else:
# This is non-windows, so we trust isatty.
OUT_TYPE = 'pipe or file'
if IS_TTY and is_windows:
# Wrapping may cause errors on some Windows versions
# (crbug.com/1114548).
if platform.release() != '10' or enable_native_ansi():
should_wrap = False
colorama.init(wrap=should_wrap)
if __name__ == '__main__':
init()
print('IS_TTY:', IS_TTY)
print('OUT_TYPE:', OUT_TYPE)