Add support for Windows 10's ANSI/VT console (#935)

Co-authored-by: Delgan <delgan.py@gmail.com>

loguru/_colorama.py

@@ -46,10 +46,21 @@ def should_wrap(stream):
 
     from colorama.win32 import winapi_test
 
-    return winapi_test()
+    if not winapi_test():
+        return False
+
+    try:
+        from colorama.winterm import enable_vt_processing
+    except ImportError:
+        return True
+
+    try:
+        return not enable_vt_processing(stream.fileno())
+    except Exception:
+        return True
 
 
 def wrap(stream):
     from colorama import AnsiToWin32
 
-    return AnsiToWin32(stream, convert=True, strip=False, autoreset=False).stream
+    return AnsiToWin32(stream, convert=True, strip=True, autoreset=False).stream

tests/conftest.py

@@ -74,21 +74,31 @@ def check_dir(dir, *, files=None, size=None):
             seen.add(filepath)
 
 
-class StreamIsattyTrue(io.StringIO):
+class StubStream(io.StringIO):
+    def fileno(self):
+        return 1
+
+
+class StreamIsattyTrue(StubStream):
     def isatty(self):
         return True
 
 
-class StreamIsattyFalse(io.StringIO):
+class StreamIsattyFalse(StubStream):
     def isatty(self):
         return False
 
 
-class StreamIsattyException(io.StringIO):
+class StreamIsattyException(StubStream):
     def isatty(self):
         raise RuntimeError
 
 
+class StreamFilenoException(StreamIsattyTrue):
+    def fileno(self):
+        raise RuntimeError
+
+
 @contextlib.contextmanager
 def default_threading_excepthook():
     if not hasattr(threading, "excepthook"):

tests/test_colorama.py

@@ -7,7 +7,12 @@
 from loguru import logger
 from loguru._colorama import should_colorize, should_wrap
 
-from .conftest import StreamIsattyException, StreamIsattyFalse, StreamIsattyTrue
+from .conftest import (
+    StreamFilenoException,
+    StreamIsattyException,
+    StreamIsattyFalse,
+    StreamIsattyTrue,
+)
 
 
 @pytest.fixture(autouse=True)
@@ -22,23 +27,38 @@ def clear_environment():
 def patch_colorama(monkeypatch):
     ansi_to_win32_class = MagicMock()
     winapi_test = MagicMock(return_value=True)
+    enable_vt_processing = MagicMock(return_value=False)
     win32 = MagicMock(winapi_test=winapi_test)
-    colorama = MagicMock(AnsiToWin32=ansi_to_win32_class, win32=win32)
+    winterm = MagicMock(enable_vt_processing=enable_vt_processing)
+    colorama = MagicMock(AnsiToWin32=ansi_to_win32_class, win32=win32, winterm=winterm)
     monkeypatch.setitem(sys.modules, "colorama", colorama)
     monkeypatch.setitem(sys.modules, "colorama.win32", win32)
+    monkeypatch.setitem(sys.modules, "colorama.winterm", winterm)
     yield colorama
 
 
 @pytest.mark.parametrize("patched", ["__stdout__", "__stderr__"])
 @pytest.mark.skipif(os.name != "nt", reason="Only Windows requires Colorama")
-def test_stream_wrapped_on_windows(patched, monkeypatch, patch_colorama):
+def test_stream_wrapped_on_windows_if_no_vt_support(patched, monkeypatch, patch_colorama):
     stream = StreamIsattyTrue()
     monkeypatch.setattr(sys, patched, stream, raising=False)
     patch_colorama.win32.winapi_test.return_value = True
+    patch_colorama.winterm.enable_vt_processing.return_value = False
     logger.add(stream, colorize=True)
     assert patch_colorama.AnsiToWin32.called
 
 
+@pytest.mark.parametrize("patched", ["__stdout__", "__stderr__"])
+@pytest.mark.skipif(os.name != "nt", reason="Only Windows requires Colorama")
+def test_stream_not_wrapped_on_windows_if_vt_support(patched, monkeypatch, patch_colorama):
+    stream = StreamIsattyTrue()
+    monkeypatch.setattr(sys, patched, stream, raising=False)
+    patch_colorama.win32.winapi_test.return_value = True
+    patch_colorama.winterm.enable_vt_processing.return_value = True
+    logger.add(stream, colorize=True)
+    assert not patch_colorama.AnsiToWin32.called
+
+
 def test_stream_is_none():
     assert not should_colorize(None)
 
@@ -170,28 +190,76 @@ def test_dont_wrap_on_linux(monkeypatch, patched, patch_colorama):
 
 @pytest.mark.parametrize("patched", ["stdout", "stderr", ""])
 @pytest.mark.skipif(os.name != "nt", reason="Only Windows requires Colorama")
-def test_dont_wrap_if_not_stdout_or_stderr(monkeypatch, patched, patch_colorama):
+def test_dont_wrap_if_not_original_stdout_or_stderr(monkeypatch, patched, patch_colorama):
     stream = StreamIsattyTrue()
     monkeypatch.setattr(sys, patched, stream, raising=False)
     assert not should_wrap(stream)
     assert not patch_colorama.win32.winapi_test.called
 
 
+@pytest.mark.parametrize("patched", ["__stdout__", "__stderr__"])
+@pytest.mark.skipif(os.name != "nt", reason="Only Windows requires Colorama")
+def test_dont_wrap_if_terminal_has_vt_support(monkeypatch, patched, patch_colorama):
+    stream = StreamIsattyTrue()
+    monkeypatch.setattr(sys, patched, stream, raising=False)
+    patch_colorama.win32.winapi_test.return_value = True
+    patch_colorama.winterm.enable_vt_processing.return_value = True
+    assert not should_wrap(stream)
+    assert patch_colorama.winterm.enable_vt_processing.called
+
+
 @pytest.mark.parametrize("patched", ["__stdout__", "__stderr__"])
 @pytest.mark.skipif(os.name != "nt", reason="Only Windows requires Colorama")
 def test_dont_wrap_if_winapi_false(monkeypatch, patched, patch_colorama):
     stream = StreamIsattyTrue()
     monkeypatch.setattr(sys, patched, stream, raising=False)
     patch_colorama.win32.winapi_test.return_value = False
+    patch_colorama.winterm.enable_vt_processing.return_value = False
     assert not should_wrap(stream)
     assert patch_colorama.win32.winapi_test.called
 
 
 @pytest.mark.parametrize("patched", ["__stdout__", "__stderr__"])
 @pytest.mark.skipif(os.name != "nt", reason="Only Windows requires Colorama")
-def test_wrap_if_winapi_true(monkeypatch, patched, patch_colorama):
+def test_wrap_if_winapi_true_and_no_vt_support(monkeypatch, patched, patch_colorama):
+    stream = StreamIsattyTrue()
+    monkeypatch.setattr(sys, patched, stream, raising=False)
+    patch_colorama.win32.winapi_test.return_value = True
+    patch_colorama.winterm.enable_vt_processing.return_value = False
+    assert should_wrap(stream)
+    assert patch_colorama.winterm.enable_vt_processing.called
+    assert patch_colorama.win32.winapi_test.called
+
+
+@pytest.mark.parametrize("patched", ["__stdout__", "__stderr__"])
+@pytest.mark.skipif(os.name != "nt", reason="Only Windows requires Colorama")
+def test_wrap_if_winapi_true_and_vt_check_fails(monkeypatch, patched, patch_colorama):
+    stream = StreamIsattyTrue()
+    monkeypatch.setattr(sys, patched, stream, raising=False)
+    patch_colorama.win32.winapi_test.return_value = True
+    patch_colorama.winterm.enable_vt_processing.side_effect = RuntimeError
+    assert should_wrap(stream)
+    assert patch_colorama.winterm.enable_vt_processing.called
+    assert patch_colorama.win32.winapi_test.called
+
+
+@pytest.mark.parametrize("patched", ["__stdout__", "__stderr__"])
+@pytest.mark.skipif(os.name != "nt", reason="Only Windows requires Colorama")
+def test_wrap_if_winapi_true_and_stream_has_no_fileno(monkeypatch, patched, patch_colorama):
+    stream = StreamFilenoException()
+    monkeypatch.setattr(sys, patched, stream, raising=False)
+    patch_colorama.win32.winapi_test.return_value = True
+    assert should_wrap(stream)
+    assert not patch_colorama.winterm.enable_vt_processing.called
+    assert patch_colorama.win32.winapi_test.called
+
+
+@pytest.mark.parametrize("patched", ["__stdout__", "__stderr__"])
+@pytest.mark.skipif(os.name != "nt", reason="Only Windows requires Colorama")
+def test_wrap_if_winapi_true_and_old_colorama_version(monkeypatch, patched, patch_colorama):
     stream = StreamIsattyTrue()
     monkeypatch.setattr(sys, patched, stream, raising=False)
     patch_colorama.win32.winapi_test.return_value = True
+    del patch_colorama.winterm.enable_vt_processing
     assert should_wrap(stream)
     assert patch_colorama.win32.winapi_test.called

tests/test_formatting.py

@@ -1,3 +1,4 @@
+import os
 import re
 
 import pytest
@@ -19,7 +20,7 @@
         ("{level.icon}", lambda r: r == "🐞"),
         ("{file}", lambda r: r == "test_formatting.py"),
         ("{file.name}", lambda r: r == "test_formatting.py"),
-        ("{file.path}", lambda r: r == __file__),
+        ("{file.path}", lambda r: os.path.normcase(r) == os.path.normcase(__file__)),
         ("{function}", lambda r: r == "test_log_formatters"),
         ("{module}", lambda r: r == "test_formatting"),
         ("{thread}", lambda r: re.fullmatch(r"\d+", r)),