Skip to content

Commit a9c6770

Browse files
committed
Greatly improved windows color support and fixed #291
1 parent 320bb54 commit a9c6770

4 files changed

Lines changed: 184 additions & 46 deletions

File tree

progressbar/env.py

Lines changed: 34 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
from __future__ import annotations
22

3+
import contextlib
34
import enum
45
import os
56
import re
@@ -8,6 +9,7 @@
89
from . import base
910

1011

12+
1113
@typing.overload
1214
def env_flag(name: str, default: bool) -> bool:
1315
...
@@ -41,6 +43,7 @@ class ColorSupport(enum.IntEnum):
4143
XTERM = 16
4244
XTERM_256 = 256
4345
XTERM_TRUECOLOR = 16777216
46+
WINDOWS = 8
4447

4548
@classmethod
4649
def from_env(cls):
@@ -65,10 +68,22 @@ def from_env(cls):
6568
)
6669

6770
if os.environ.get('JUPYTER_COLUMNS') or os.environ.get(
68-
'JUPYTER_LINES',
71+
'JUPYTER_LINES',
6972
):
7073
# Jupyter notebook always supports true color.
7174
return cls.XTERM_TRUECOLOR
75+
elif os.name == 'nt':
76+
# We can't reliably detect true color support on Windows, so we
77+
# will assume it is supported if the console is configured to
78+
# support it.
79+
from .terminal.os_specific import windows
80+
if (
81+
windows.get_console_mode() &
82+
windows.WindowsConsoleModeFlags.ENABLE_PROCESSED_OUTPUT
83+
):
84+
return cls.XTERM_TRUECOLOR
85+
else:
86+
return cls.WINDOWS
7287

7388
support = cls.NONE
7489
for variable in variables:
@@ -88,17 +103,17 @@ def from_env(cls):
88103

89104

90105
def is_ansi_terminal(
91-
fd: base.IO,
92-
is_terminal: bool | None = None,
93-
) -> bool: # pragma: no cover
106+
fd: base.IO,
107+
is_terminal: bool | None = None,
108+
) -> bool | None: # pragma: no cover
94109
if is_terminal is None:
95110
# Jupyter Notebooks define this variable and support progress bars
96111
if 'JPY_PARENT_PID' in os.environ:
97112
is_terminal = True
98113
# This works for newer versions of pycharm only. With older versions
99114
# there is no way to check.
100115
elif os.environ.get('PYCHARM_HOSTED') == '1' and not os.environ.get(
101-
'PYTEST_CURRENT_TEST',
116+
'PYTEST_CURRENT_TEST',
102117
):
103118
is_terminal = True
104119

@@ -108,20 +123,24 @@ def is_ansi_terminal(
108123
# isatty has not been defined we have no way of knowing so we will not
109124
# use ansi. ansi terminals will typically define one of the 2
110125
# environment variables.
111-
try:
126+
with contextlib.suppress(Exception):
112127
is_tty = fd.isatty()
113128
# Try and match any of the huge amount of Linux/Unix ANSI consoles
114129
if is_tty and ANSI_TERM_RE.match(os.environ.get('TERM', '')):
115130
is_terminal = True
116131
# ANSICON is a Windows ANSI compatible console
117132
elif 'ANSICON' in os.environ:
118133
is_terminal = True
134+
elif os.name == 'nt':
135+
from .terminal.os_specific import windows
136+
return bool(
137+
windows.get_console_mode() &
138+
windows.WindowsConsoleModeFlags.ENABLE_PROCESSED_OUTPUT
139+
)
119140
else:
120141
is_terminal = None
121-
except Exception:
122-
is_terminal = False
123142

124-
return bool(is_terminal)
143+
return is_terminal
125144

126145

127146
def is_terminal(fd: base.IO, is_terminal: bool | None = None) -> bool:
@@ -144,6 +163,12 @@ def is_terminal(fd: base.IO, is_terminal: bool | None = None) -> bool:
144163
return bool(is_terminal)
145164

146165

166+
# Enable Windows full color mode if possible
167+
if os.name == 'nt':
168+
from .terminal import os_specific
169+
170+
os_specific.set_console_mode()
171+
147172
COLOR_SUPPORT = ColorSupport.from_env()
148173
ANSI_TERMS = (
149174
'([xe]|bv)term',

progressbar/terminal/base.py

Lines changed: 98 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -3,20 +3,20 @@
33
import abc
44
import collections
55
import colorsys
6+
import enum
67
import threading
78
from collections import defaultdict
8-
99
# Ruff is being stupid and doesn't understand `ClassVar` if it comes from the
1010
# `types` module
1111
from typing import ClassVar
1212

1313
from python_utils import converters, types
1414

15+
from .os_specific import getch
1516
from .. import (
1617
base as pbase,
1718
env,
1819
)
19-
from .os_specific import getch
2020

2121
ESC = '\x1B'
2222

@@ -178,6 +178,53 @@ def column(self, stream):
178178
return column
179179

180180

181+
182+
class WindowsColors(enum.Enum):
183+
BLACK = 0, 0, 0
184+
BLUE = 0, 0, 128
185+
GREEN = 0, 128, 0
186+
CYAN = 0, 128, 128
187+
RED = 128, 0, 0
188+
MAGENTA = 128, 0, 128
189+
YELLOW = 128, 128, 0
190+
GREY = 192, 192, 192
191+
INTENSE_BLACK = 128, 128, 128
192+
INTENSE_BLUE = 0, 0, 255
193+
INTENSE_GREEN = 0, 255, 0
194+
INTENSE_CYAN = 0, 255, 255
195+
INTENSE_RED = 255, 0, 0
196+
INTENSE_MAGENTA = 255, 0, 255
197+
INTENSE_YELLOW = 255, 255, 0
198+
INTENSE_WHITE = 255, 255, 255
199+
200+
@staticmethod
201+
def from_rgb(rgb: types.Tuple[int, int, int]):
202+
"""Find the closest ConsoleColor to the given RGB color."""
203+
204+
def color_distance(rgb1, rgb2):
205+
return sum((c1 - c2) ** 2 for c1, c2 in zip(rgb1, rgb2))
206+
207+
return min(
208+
WindowsColors,
209+
key=lambda color: color_distance(color.value, rgb),
210+
)
211+
212+
213+
class WindowsColor:
214+
__slots__ = 'color',
215+
216+
def __init__(self, color: Color):
217+
self.color = color
218+
219+
def __call__(self, text):
220+
return text
221+
# In the future we might want to use this, but it requires direct printing to stdout and all of our surrounding functions expect buffered output so it's not feasible right now.
222+
# Additionally, recent Windows versions all support ANSI codes without issue so there is little need.
223+
# from progressbar.terminal.os_specific import windows
224+
# windows.print_color(text, WindowsColors.from_rgb(self.color.rgb))
225+
226+
227+
181228
class RGB(collections.namedtuple('RGB', ['red', 'green', 'blue'])):
182229
__slots__ = ()
183230

@@ -207,6 +254,14 @@ def to_ansi_256(self):
207254
blue = round(self.blue / 255 * 5)
208255
return 16 + 36 * red + 6 * green + blue
209256

257+
@property
258+
def to_windows(self):
259+
'''
260+
Convert an RGB color (0-255 per channel) to the closest color in the
261+
Windows 16 color scheme.
262+
'''
263+
return WindowsColors.from_rgb((self.red, self.green, self.blue))
264+
210265
def interpolate(self, end: RGB, step: float) -> RGB:
211266
return RGB(
212267
int(self.red + (end.red - self.red) * step),
@@ -286,27 +341,36 @@ def __call__(self, value: str) -> str:
286341

287342
@property
288343
def fg(self):
289-
return SGRColor(self, 38, 39)
344+
if env.COLOR_SUPPORT is env.ColorSupport.WINDOWS:
345+
return WindowsColor(self)
346+
else:
347+
return SGRColor(self, 38, 39)
290348

291349
@property
292350
def bg(self):
293-
return SGRColor(self, 48, 49)
351+
if env.COLOR_SUPPORT is env.ColorSupport.WINDOWS:
352+
return DummyColor()
353+
else:
354+
return SGRColor(self, 48, 49)
294355

295356
@property
296357
def underline(self):
297-
return SGRColor(self, 58, 59)
358+
if env.COLOR_SUPPORT is env.ColorSupport.WINDOWS:
359+
return DummyColor()
360+
else:
361+
return SGRColor(self, 58, 59)
298362

299363
@property
300364
def ansi(self) -> types.Optional[str]:
301365
if (
302-
env.COLOR_SUPPORT is env.ColorSupport.XTERM_TRUECOLOR
366+
env.COLOR_SUPPORT is env.ColorSupport.XTERM_TRUECOLOR
303367
): # pragma: no branch
304368
return f'2;{self.rgb.red};{self.rgb.green};{self.rgb.blue}'
305369

306370
if self.xterm: # pragma: no branch
307371
color = self.xterm
308372
elif (
309-
env.COLOR_SUPPORT is env.ColorSupport.XTERM_256
373+
env.COLOR_SUPPORT is env.ColorSupport.XTERM_256
310374
): # pragma: no branch
311375
color = self.rgb.to_ansi_256
312376
elif env.COLOR_SUPPORT is env.ColorSupport.XTERM: # pragma: no branch
@@ -354,11 +418,11 @@ class Colors:
354418

355419
@classmethod
356420
def register(
357-
cls,
358-
rgb: RGB,
359-
hls: types.Optional[HSL] = None,
360-
name: types.Optional[str] = None,
361-
xterm: types.Optional[int] = None,
421+
cls,
422+
rgb: RGB,
423+
hls: types.Optional[HSL] = None,
424+
name: types.Optional[str] = None,
425+
xterm: types.Optional[int] = None,
362426
) -> Color:
363427
color = Color(rgb, hls, name, xterm)
364428

@@ -395,9 +459,9 @@ def __call__(self, value: float) -> Color:
395459
def get_color(self, value: float) -> Color:
396460
'Map a value from 0 to 1 to a color.'
397461
if (
398-
value == pbase.Undefined
399-
or value == pbase.UnknownLength
400-
or value <= 0
462+
value == pbase.Undefined
463+
or value == pbase.UnknownLength
464+
or value <= 0
401465
):
402466
return self.colors[0]
403467
elif value >= 1:
@@ -443,14 +507,14 @@ def get_color(value: float, color: OptionalColor) -> Color | None:
443507

444508

445509
def apply_colors(
446-
text: str,
447-
percentage: float | None = None,
448-
*,
449-
fg: OptionalColor = None,
450-
bg: OptionalColor = None,
451-
fg_none: Color | None = None,
452-
bg_none: Color | None = None,
453-
**kwargs: types.Any,
510+
text: str,
511+
percentage: float | None = None,
512+
*,
513+
fg: OptionalColor = None,
514+
bg: OptionalColor = None,
515+
fg_none: Color | None = None,
516+
bg_none: Color | None = None,
517+
**kwargs: types.Any,
454518
) -> str:
455519
'''Apply colors/gradients to a string depending on the given percentage.
456520
@@ -475,6 +539,17 @@ def apply_colors(
475539
return text
476540

477541

542+
class DummyColor:
543+
def __call__(self, text):
544+
return text
545+
546+
def __getattr__(self, item):
547+
return self
548+
549+
def __repr__(self):
550+
return 'DummyColor()'
551+
552+
478553
class SGR(CSI):
479554
_start_code: int
480555
_end_code: int

progressbar/terminal/os_specific/__init__.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
getch as _getch,
66
reset_console_mode as _reset_console_mode,
77
set_console_mode as _set_console_mode,
8+
get_console_mode as _get_console_mode,
89
)
910

1011
else:
@@ -16,7 +17,11 @@ def _reset_console_mode():
1617
def _set_console_mode():
1718
pass
1819

20+
def _get_console_mode():
21+
return 0
22+
1923

2024
getch = _getch
2125
reset_console_mode = _reset_console_mode
2226
set_console_mode = _set_console_mode
27+
get_console_mode = _get_console_mode

0 commit comments

Comments
 (0)