Greatly improved windows color support and fixed wolph#291

wolph · wolph · commit a9c677021e96 · 2024-01-25T03:49:22.000+01:00
diff --git a/progressbar/env.py b/progressbar/env.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import contextlib
 import enum
 import os
 import re
@@ -8,6 +9,7 @@
 from . import base
 
 
+
 @typing.overload
 def env_flag(name: str, default: bool) -> bool:
     ...
@@ -41,6 +43,7 @@ class ColorSupport(enum.IntEnum):
     XTERM = 16
     XTERM_256 = 256
     XTERM_TRUECOLOR = 16777216
+    WINDOWS = 8
 
     @classmethod
     def from_env(cls):
@@ -65,10 +68,22 @@ def from_env(cls):
         )
 
         if os.environ.get('JUPYTER_COLUMNS') or os.environ.get(
-            'JUPYTER_LINES',
+                'JUPYTER_LINES',
         ):
             # Jupyter notebook always supports true color.
             return cls.XTERM_TRUECOLOR
+        elif os.name == 'nt':
+            # We can't reliably detect true color support on Windows, so we
+            # will assume it is supported if the console is configured to
+            # support it.
+            from .terminal.os_specific import windows
+            if (
+                    windows.get_console_mode() &
+                    windows.WindowsConsoleModeFlags.ENABLE_PROCESSED_OUTPUT
+            ):
+                return cls.XTERM_TRUECOLOR
+            else:
+                return cls.WINDOWS
 
         support = cls.NONE
         for variable in variables:
@@ -88,17 +103,17 @@ def from_env(cls):
 
 
 def is_ansi_terminal(
-    fd: base.IO,
-    is_terminal: bool | None = None,
-) -> bool:  # pragma: no cover
+        fd: base.IO,
+        is_terminal: bool | None = None,
+) -> bool | None:  # pragma: no cover
     if is_terminal is None:
         # Jupyter Notebooks define this variable and support progress bars
         if 'JPY_PARENT_PID' in os.environ:
             is_terminal = True
         # This works for newer versions of pycharm only. With older versions
         # there is no way to check.
         elif os.environ.get('PYCHARM_HOSTED') == '1' and not os.environ.get(
-            'PYTEST_CURRENT_TEST',
+                'PYTEST_CURRENT_TEST',
         ):
             is_terminal = True
 
@@ -108,20 +123,24 @@ def is_ansi_terminal(
         # isatty has not been defined we have no way of knowing so we will not
         # use ansi.  ansi terminals will typically define one of the 2
         # environment variables.
-        try:
+        with contextlib.suppress(Exception):
             is_tty = fd.isatty()
             # Try and match any of the huge amount of Linux/Unix ANSI consoles
             if is_tty and ANSI_TERM_RE.match(os.environ.get('TERM', '')):
                 is_terminal = True
             # ANSICON is a Windows ANSI compatible console
             elif 'ANSICON' in os.environ:
                 is_terminal = True
+            elif os.name == 'nt':
+                from .terminal.os_specific import windows
+                return bool(
+                    windows.get_console_mode() &
+                    windows.WindowsConsoleModeFlags.ENABLE_PROCESSED_OUTPUT
+                )
             else:
                 is_terminal = None
-        except Exception:
-            is_terminal = False
 
-    return bool(is_terminal)
+    return is_terminal
 
 
 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:
     return bool(is_terminal)
 
 
+# Enable Windows full color mode if possible
+if os.name == 'nt':
+    from .terminal import os_specific
+
+    os_specific.set_console_mode()
+
 COLOR_SUPPORT = ColorSupport.from_env()
 ANSI_TERMS = (
     '([xe]|bv)term',
diff --git a/progressbar/terminal/base.py b/progressbar/terminal/base.py
@@ -3,20 +3,20 @@
 import abc
 import collections
 import colorsys
+import enum
 import threading
 from collections import defaultdict
-
 # Ruff is being stupid and doesn't understand `ClassVar` if it comes from the
 # `types` module
 from typing import ClassVar
 
 from python_utils import converters, types
 
+from .os_specific import getch
 from .. import (
     base as pbase,
     env,
 )
-from .os_specific import getch
 
 ESC = '\x1B'
 
@@ -178,6 +178,53 @@ def column(self, stream):
         return column
 
 
+
+class WindowsColors(enum.Enum):
+    BLACK = 0, 0, 0
+    BLUE = 0, 0, 128
+    GREEN = 0, 128, 0
+    CYAN = 0, 128, 128
+    RED = 128, 0, 0
+    MAGENTA = 128, 0, 128
+    YELLOW = 128, 128, 0
+    GREY = 192, 192, 192
+    INTENSE_BLACK = 128, 128, 128
+    INTENSE_BLUE = 0, 0, 255
+    INTENSE_GREEN = 0, 255, 0
+    INTENSE_CYAN = 0, 255, 255
+    INTENSE_RED = 255, 0, 0
+    INTENSE_MAGENTA = 255, 0, 255
+    INTENSE_YELLOW = 255, 255, 0
+    INTENSE_WHITE = 255, 255, 255
+
+    @staticmethod
+    def from_rgb(rgb: types.Tuple[int, int, int]):
+        """Find the closest ConsoleColor to the given RGB color."""
+
+        def color_distance(rgb1, rgb2):
+            return sum((c1 - c2) ** 2 for c1, c2 in zip(rgb1, rgb2))
+
+        return min(
+            WindowsColors,
+            key=lambda color: color_distance(color.value, rgb),
+        )
+
+
+class WindowsColor:
+    __slots__ = 'color',
+
+    def __init__(self, color: Color):
+        self.color = color
+
+    def __call__(self, text):
+        return text
+        # 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.
+        # Additionally, recent Windows versions all support ANSI codes without issue so there is little need.
+        # from progressbar.terminal.os_specific import windows
+        # windows.print_color(text, WindowsColors.from_rgb(self.color.rgb))
+
+
+
 class RGB(collections.namedtuple('RGB', ['red', 'green', 'blue'])):
     __slots__ = ()
 
@@ -207,6 +254,14 @@ def to_ansi_256(self):
         blue = round(self.blue / 255 * 5)
         return 16 + 36 * red + 6 * green + blue
 
+    @property
+    def to_windows(self):
+        '''
+        Convert an RGB color (0-255 per channel) to the closest color in the
+        Windows 16 color scheme.
+        '''
+        return WindowsColors.from_rgb((self.red, self.green, self.blue))
+
     def interpolate(self, end: RGB, step: float) -> RGB:
         return RGB(
             int(self.red + (end.red - self.red) * step),
@@ -286,27 +341,36 @@ def __call__(self, value: str) -> str:
 
     @property
     def fg(self):
-        return SGRColor(self, 38, 39)
+        if env.COLOR_SUPPORT is env.ColorSupport.WINDOWS:
+            return WindowsColor(self)
+        else:
+            return SGRColor(self, 38, 39)
 
     @property
     def bg(self):
-        return SGRColor(self, 48, 49)
+        if env.COLOR_SUPPORT is env.ColorSupport.WINDOWS:
+            return DummyColor()
+        else:
+            return SGRColor(self, 48, 49)
 
     @property
     def underline(self):
-        return SGRColor(self, 58, 59)
+        if env.COLOR_SUPPORT is env.ColorSupport.WINDOWS:
+            return DummyColor()
+        else:
+            return SGRColor(self, 58, 59)
 
     @property
     def ansi(self) -> types.Optional[str]:
         if (
-            env.COLOR_SUPPORT is env.ColorSupport.XTERM_TRUECOLOR
+                env.COLOR_SUPPORT is env.ColorSupport.XTERM_TRUECOLOR
         ):  # pragma: no branch
             return f'2;{self.rgb.red};{self.rgb.green};{self.rgb.blue}'
 
         if self.xterm:  # pragma: no branch
             color = self.xterm
         elif (
-            env.COLOR_SUPPORT is env.ColorSupport.XTERM_256
+                env.COLOR_SUPPORT is env.ColorSupport.XTERM_256
         ):  # pragma: no branch
             color = self.rgb.to_ansi_256
         elif env.COLOR_SUPPORT is env.ColorSupport.XTERM:  # pragma: no branch
@@ -354,11 +418,11 @@ class Colors:
 
     @classmethod
     def register(
-        cls,
-        rgb: RGB,
-        hls: types.Optional[HSL] = None,
-        name: types.Optional[str] = None,
-        xterm: types.Optional[int] = None,
+            cls,
+            rgb: RGB,
+            hls: types.Optional[HSL] = None,
+            name: types.Optional[str] = None,
+            xterm: types.Optional[int] = None,
     ) -> Color:
         color = Color(rgb, hls, name, xterm)
 
@@ -395,9 +459,9 @@ def __call__(self, value: float) -> Color:
     def get_color(self, value: float) -> Color:
         'Map a value from 0 to 1 to a color.'
         if (
-            value == pbase.Undefined
-            or value == pbase.UnknownLength
-            or value <= 0
+                value == pbase.Undefined
+                or value == pbase.UnknownLength
+                or value <= 0
         ):
             return self.colors[0]
         elif value >= 1:
@@ -443,14 +507,14 @@ def get_color(value: float, color: OptionalColor) -> Color | None:
 
 
 def apply_colors(
-    text: str,
-    percentage: float | None = None,
-    *,
-    fg: OptionalColor = None,
-    bg: OptionalColor = None,
-    fg_none: Color | None = None,
-    bg_none: Color | None = None,
-    **kwargs: types.Any,
+        text: str,
+        percentage: float | None = None,
+        *,
+        fg: OptionalColor = None,
+        bg: OptionalColor = None,
+        fg_none: Color | None = None,
+        bg_none: Color | None = None,
+        **kwargs: types.Any,
 ) -> str:
     '''Apply colors/gradients to a string depending on the given percentage.
 
@@ -475,6 +539,17 @@ def apply_colors(
     return text
 
 
+class DummyColor:
+    def __call__(self, text):
+        return text
+
+    def __getattr__(self, item):
+        return self
+
+    def __repr__(self):
+        return 'DummyColor()'
+
+
 class SGR(CSI):
     _start_code: int
     _end_code: int
diff --git a/progressbar/terminal/os_specific/__init__.py b/progressbar/terminal/os_specific/__init__.py
@@ -5,6 +5,7 @@
         getch as _getch,
         reset_console_mode as _reset_console_mode,
         set_console_mode as _set_console_mode,
+        get_console_mode as _get_console_mode,
     )
 
 else:
@@ -16,7 +17,11 @@ def _reset_console_mode():
     def _set_console_mode():
         pass
 
+    def _get_console_mode():
+        return 0
+
 
 getch = _getch
 reset_console_mode = _reset_console_mode
 set_console_mode = _set_console_mode
+get_console_mode = _get_console_mode
diff --git a/progressbar/terminal/os_specific/windows.py b/progressbar/terminal/os_specific/windows.py
