1111import logging
1212
1313
14- def add_coloring_to_emit_windows (fn ):
14+ def ansi_color_available ():
15+ if not sys .platform .startswith ('win' ):
16+ return sys .stderr .isatty ()
17+
1518 # Constants from the Windows API
1619 STD_OUTPUT_HANDLE = - 11
20+ ENABLE_VIRTUAL_TERMINAL_PROCESSING = 0x0004
1721
18- def _get_color ():
19- SHORT = ctypes .c_short
20- WORD = ctypes .c_ushort
21-
22- class COORD (ctypes .Structure ):
23- _fields_ = [
24- ("X" , SHORT ),
25- ("Y" , SHORT )]
26-
27- class SMALL_RECT (ctypes .Structure ):
28- _fields_ = [
29- ("Left" , SHORT ),
30- ("Top" , SHORT ),
31- ("Right" , SHORT ),
32- ("Bottom" , SHORT )]
33-
34- class CONSOLE_SCREEN_BUFFER_INFO (ctypes .Structure ):
35- _fields_ = [
36- ("dwSize" , COORD ),
37- ("dwCursorPosition" , COORD ),
38- ("wAttributes" , WORD ),
39- ("srWindow" , SMALL_RECT ),
40- ("dwMaximumWindowSize" , COORD )]
22+ kernel32 = ctypes .windll .kernel32
23+ stdout_handle = kernel32 .GetStdHandle (STD_OUTPUT_HANDLE )
4124
42- hdl = ctypes .windll .kernel32 .GetStdHandle (STD_OUTPUT_HANDLE )
43- csbi = CONSOLE_SCREEN_BUFFER_INFO ()
44- ctypes .windll .kernel32 .GetConsoleScreenBufferInfo (hdl , ctypes .byref (csbi ))
45- return csbi .wAttributes
25+ # Get the current console mode
26+ console_mode = ctypes .c_uint ()
27+ if not kernel32 .GetConsoleMode (stdout_handle , ctypes .byref (console_mode )):
28+ # Handle error if GetConsoleMode fails
29+ return False
4630
47- def _set_color (code ):
48- hdl = ctypes .windll .kernel32 .GetStdHandle (STD_OUTPUT_HANDLE )
49- ctypes .windll .kernel32 .SetConsoleTextAttribute (hdl , code )
50-
51- def new (* args ):
52- # wincon.h
53- FOREGROUND_BLACK = 0x0000 # noqa
54- FOREGROUND_BLUE = 0x0001 # noqa
55- FOREGROUND_GREEN = 0x0002 # noqa
56- FOREGROUND_CYAN = 0x0003 # noqa
57- FOREGROUND_RED = 0x0004 # noqa
58- FOREGROUND_MAGENTA = 0x0005 # noqa
59- FOREGROUND_YELLOW = 0x0006 # noqa
60- FOREGROUND_GREY = 0x0007 # noqa
61- FOREGROUND_INTENSITY = 0x0008 # foreground color is intensified.
62-
63- FOREGROUND_WHITE = FOREGROUND_BLUE | FOREGROUND_GREEN | FOREGROUND_RED # noqa
64-
65- BACKGROUND_BLACK = 0x0000 # noqa
66- BACKGROUND_BLUE = 0x0010 # noqa
67- BACKGROUND_GREEN = 0x0020 # noqa
68- BACKGROUND_CYAN = 0x0030 # noqa
69- BACKGROUND_RED = 0x0040 # noqa
70- BACKGROUND_MAGENTA = 0x0050 # noqa
71- BACKGROUND_YELLOW = 0x0060 # noqa
72- BACKGROUND_GREY = 0x0070 # noqa
73- BACKGROUND_INTENSITY = 0x0080 # background color is intensified.
74- levelno = args [1 ].levelno
75- if (levelno >= 50 ):
76- color = BACKGROUND_YELLOW | FOREGROUND_RED | FOREGROUND_INTENSITY | BACKGROUND_INTENSITY
77- elif (levelno >= 40 ):
78- color = FOREGROUND_RED | FOREGROUND_INTENSITY
79- elif (levelno >= 30 ):
80- color = FOREGROUND_YELLOW | FOREGROUND_INTENSITY
81- elif (levelno >= 20 ):
82- color = FOREGROUND_GREEN
83- elif (levelno >= 10 ):
84- color = FOREGROUND_MAGENTA
85- else :
86- color = FOREGROUND_WHITE
87-
88- old_color = _get_color ()
89- _set_color (color )
90- ret = fn (* args )
91- _set_color (old_color )
92- return ret
93-
94- new .orig_func = fn
95- return new
31+ # Check if the flag is set in the current console mode
32+ return (console_mode .value & ENABLE_VIRTUAL_TERMINAL_PROCESSING ) != 0
9633
9734
9835def add_coloring_to_emit_ansi (fn ):
@@ -117,11 +54,8 @@ def new(*args):
11754
11855
11956def enable ():
120- if sys .stderr .isatty ():
121- if sys .platform .startswith ('win' ):
122- logging .StreamHandler .emit = add_coloring_to_emit_windows (logging .StreamHandler .emit )
123- else :
124- logging .StreamHandler .emit = add_coloring_to_emit_ansi (logging .StreamHandler .emit )
57+ if ansi_color_available ():
58+ logging .StreamHandler .emit = add_coloring_to_emit_ansi (logging .StreamHandler .emit )
12559
12660
12761def disable ():
0 commit comments