Add colorama.just_fix_windows_console()

README.rst

@@ -55,7 +55,8 @@ This has the upshot of providing a simple cross-platform API for printing
 colored terminal text from Python, and has the happy side-effect that existing
 applications or libraries which use ANSI sequences to produce colored output on
 Linux or Macs can now also work on Windows, simply by calling
-``colorama.init()``.
+``colorama.just_fix_windows_console()`` (since v0.4.6) or ``colorama.init()``
+(all versions, but may have other side-effects – see below).
 
 An alternative approach is to install ``ansi.sys`` on Windows machines, which
 provides the same behaviour for all applications running in terminals. Colorama
@@ -85,30 +86,66 @@ Usage
 Initialisation
 ..............
 
-Applications should initialise Colorama using:
+If the only thing you want from Colorama is to get ANSI escapes to work on
+Windows, then run:
+
+.. code-block:: python
+
+    from colorama import just_fix_windows_console
+    just_fix_windows_console()
+
+If you're on a recent version of Windows 10 or better, and your stdout/stderr
+are pointing to a Windows console, then this will flip the magic configuration
+switch to enable Windows' built-in ANSI support.
+
+If you're on an older version of Windows, and your stdout/stderr are pointing to
+a Windows console, then this will wrap ``sys.stdout`` and/or ``sys.stderr`` in a
+magic file object that intercepts ANSI escape sequences and issues the
+appropriate Win32 calls to emulate them.
+
+In all other circumstances, it does nothing whatsoever. Basically the idea is
+that this makes Windows act like Unix with respect to ANSI escape handling.
+
+It's safe to call this function multiple times. It's safe to call this function
+on non-Windows platforms, but it won't do anything. It's safe to call this
+function when one or both of your stdout/stderr are redirected to a file – it
+won't do anything to those streams.
+
+Alternatively, you can use the older interface with more features (but also more
+potential footguns):
 
 .. code-block:: python
 
     from colorama import init
     init()
 
-On Windows, calling ``init()`` will filter ANSI escape sequences out of any
-text sent to ``stdout`` or ``stderr``, and replace them with equivalent Win32
-calls.
+This does the same thing as ``just_fix_windows_console``, except for the
+following differences:
+
+- It's not safe to call ``init`` multiple times; you can end up with multiple
+  layers of wrapping and broken ANSI support.
 
-On other platforms, calling ``init()`` has no effect (unless you request other
-optional functionality, see "Init Keyword Args" below; or if output
-is redirected). By design, this permits applications to call ``init()``
-unconditionally on all platforms, after which ANSI output should just work.
+- Colorama will apply a heuristic to guess whether stdout/stderr support ANSI,
+  and if it thinks they don't, then it will wrap ``sys.stdout`` and
+  ``sys.stderr`` in a magic file object that strips out ANSI escape sequences
+  before printing them. This happens on all platforms, and can be convenient if
+  you want to write your code to emit ANSI escape sequences unconditionally, and
+  let Colorama decide whether they should actually be output. But note that
+  Colorama's heuristic is not particularly clever.
 
-On all platforms, if output is redirected, ANSI escape sequences are completely
-stripped out.
+- ``init`` also accepts explicit keyword args to enable/disable various
+  functionality – see below.
 
 To stop using Colorama before your program exits, simply call ``deinit()``.
 This will restore ``stdout`` and ``stderr`` to their original values, so that
 Colorama is disabled. To resume using Colorama again, call ``reinit()``; it is
 cheaper than calling ``init()`` again (but does the same thing).
 
+Most users should depend on ``colorama >= 0.4.6``, and use
+``just_fix_windows_console``. The old ``init`` interface will be supported
+indefinitely for backwards compatibility, but we don't plan to fix any issues
+with it, also for backwards compatibility.
+
 
 Colored Output
 ..............
@@ -145,11 +182,11 @@ those ANSI sequences to also work on Windows:
 
 .. code-block:: python
 
-    from colorama import init
+    from colorama import just_fix_windows_console
     from termcolor import colored
 
     # use Colorama to make Termcolor work on Windows too
-    init()
+    just_fix_windows_console()
 
     # then use Termcolor for all colored text output
     print(colored('Hello, World!', 'green', 'on_red'))

colorama/__init__.py

@@ -1,5 +1,5 @@
 # Copyright Jonathan Hartley 2013. BSD 3-Clause license, see LICENSE file.
-from .initialise import init, deinit, reinit, colorama_text
+from .initialise import init, deinit, reinit, colorama_text, just_fix_windows_console
 from .ansi import Fore, Back, Style, Cursor
 from .ansitowin32 import AnsiToWin32
 

colorama/ansitowin32.py

@@ -271,3 +271,7 @@ def convert_osc(self, text):
                     if params[0] in '02':
                         winterm.set_title(params[1])
         return text
+
+
+    def flush(self):
+        self.wrapped.flush()

colorama/initialise.py

@@ -6,13 +6,27 @@
 from .ansitowin32 import AnsiToWin32
 
 
-orig_stdout = None
-orig_stderr = None
+def _wipe_internal_state_for_tests():
+    global orig_stdout, orig_stderr
+    orig_stdout = None
+    orig_stderr = None
+
+    global wrapped_stdout, wrapped_stderr
+    wrapped_stdout = None
+    wrapped_stderr = None
 
-wrapped_stdout = None
-wrapped_stderr = None
+    global atexit_done
+    atexit_done = False
+
+    global fixed_windows_console
+    fixed_windows_console = False
 
-atexit_done = False
+    try:
+        # no-op if it wasn't registered
+        atexit.unregister(reset_all)
+    except AttributeError:
+        # python 2: no atexit.unregister. Oh well, we did our best.
+        pass
 
 
 def reset_all():
@@ -55,6 +69,29 @@ def deinit():
         sys.stderr = orig_stderr
 
 
+def just_fix_windows_console():
+    global fixed_windows_console
+
+    if sys.platform != "win32":
+        return
+    if fixed_windows_console:
+        return
+    if wrapped_stdout is not None or wrapped_stderr is not None:
+        # Someone already ran init() and it did stuff, so we won't second-guess them
+        return
+
+    # On newer versions of Windows, AnsiToWin32.__init__ will implicitly enable the
+    # native ANSI support in the console as a side-effect. We only need to actually
+    # replace sys.stdout/stderr if we're in the old-style conversion mode.
+    new_stdout = AnsiToWin32(sys.stdout, convert=None, strip=None, autoreset=False)
+    if new_stdout.convert:
+        sys.stdout = new_stdout
+    new_stderr = AnsiToWin32(sys.stderr, convert=None, strip=None, autoreset=False)
+    if new_stderr.convert:
+        sys.stderr = new_stderr
+
+    fixed_windows_console = True
+
 @contextlib.contextmanager
 def colorama_text(*args, **kwargs):
     init(*args, **kwargs)
@@ -78,3 +115,7 @@ def wrap_stream(stream, convert, strip, autoreset, wrap):
         if wrapper.should_wrap():
             stream = wrapper.stream
     return stream
+
+
+# Use this for initial setup as well, to reduce code duplication
+_wipe_internal_state_for_tests()

colorama/tests/initialise_test.py

@@ -3,12 +3,12 @@
 from unittest import TestCase, main, skipUnless
 
 try:
-    from unittest.mock import patch
+    from unittest.mock import patch, Mock
 except ImportError:
-    from mock import patch
+    from mock import patch, Mock
 
 from ..ansitowin32 import StreamWrapper
-from ..initialise import init
+from ..initialise import init, just_fix_windows_console, _wipe_internal_state_for_tests
 from .utils import osname, replace_by
 
 orig_stdout = sys.stdout
@@ -23,6 +23,7 @@ def setUp(self):
         self.assertNotWrapped()
 
     def tearDown(self):
+        _wipe_internal_state_for_tests()
         sys.stdout = orig_stdout
         sys.stderr = orig_stderr
 
@@ -40,6 +41,7 @@ def assertNotWrapped(self):
 
     @patch('colorama.initialise.reset_all')
     @patch('colorama.ansitowin32.winapi_test', lambda *_: True)
+    @patch('colorama.ansitowin32.enable_vt_processing', lambda *_: False)
     def testInitWrapsOnWindows(self, _):
         with osname("nt"):
             init()
@@ -78,14 +80,6 @@ def testInitWrapOffDoesntWrapOnWindows(self):
     def testInitWrapOffIncompatibleWithAutoresetOn(self):
         self.assertRaises(ValueError, lambda: init(autoreset=True, wrap=False))
 
-    @patch('colorama.ansitowin32.winterm', None)
-    @patch('colorama.ansitowin32.winapi_test', lambda *_: True)
-    def testInitOnlyWrapsOnce(self):
-        with osname("nt"):
-            init()
-            init()
-            self.assertWrapped()
-
     @patch('colorama.win32.SetConsoleTextAttribute')
     @patch('colorama.initialise.AnsiToWin32')
     def testAutoResetPassedOn(self, mockATW32, _):
@@ -122,5 +116,74 @@ def testAtexitRegisteredOnlyOnce(self, mockRegister):
         self.assertFalse(mockRegister.called)
 
 
+class JustFixWindowsConsoleTest(TestCase):
+    def _reset(self):
+        _wipe_internal_state_for_tests()
+        sys.stdout = orig_stdout
+        sys.stderr = orig_stderr
+
+    def tearDown(self):
+        self._reset()
+
+    @patch("colorama.ansitowin32.winapi_test", lambda: True)
+    def testJustFixWindowsConsole(self):
+        if sys.platform != "win32":
+            # just_fix_windows_console should be a no-op
+            just_fix_windows_console()
+            self.assertIs(sys.stdout, orig_stdout)
+            self.assertIs(sys.stderr, orig_stderr)
+        else:
+            def fake_std():
+                # Emulate stdout=not a tty, stderr=tty
+                # to check that we handle both cases correctly
+                stdout = Mock()
+                stdout.closed = False
+                stdout.isatty.return_value = False
+                stdout.fileno.return_value = 1
+                sys.stdout = stdout
+
+                stderr = Mock()
+                stderr.closed = False
+                stderr.isatty.return_value = True
+                stderr.fileno.return_value = 2
+                sys.stderr = stderr
+
+            for native_ansi in [False, True]:
+                with patch(
+                    'colorama.ansitowin32.enable_vt_processing',
+                    lambda *_: native_ansi
+                ):
+                    self._reset()
+                    fake_std()
+
+                    # Regular single-call test
+                    prev_stdout = sys.stdout
+                    prev_stderr = sys.stderr
+                    just_fix_windows_console()
+                    self.assertIs(sys.stdout, prev_stdout)
+                    if native_ansi:
+                        self.assertIs(sys.stderr, prev_stderr)
+                    else:
+                        self.assertIsNot(sys.stderr, prev_stderr)
+
+                    # second call without resetting is always a no-op
+                    prev_stdout = sys.stdout
+                    prev_stderr = sys.stderr
+                    just_fix_windows_console()
+                    self.assertIs(sys.stdout, prev_stdout)
+                    self.assertIs(sys.stderr, prev_stderr)
+
+                    self._reset()
+                    fake_std()
+
+                    # If init() runs first, just_fix_windows_console should be a no-op
+                    init()
+                    prev_stdout = sys.stdout
+                    prev_stderr = sys.stderr
+                    just_fix_windows_console()
+                    self.assertIs(prev_stdout, sys.stdout)
+                    self.assertIs(prev_stderr, sys.stderr)
+
+
 if __name__ == '__main__':
     main()

colorama/winterm.py

@@ -190,5 +190,6 @@ def enable_vt_processing(fd):
         mode = win32.GetConsoleMode(handle)
         if mode & win32.ENABLE_VIRTUAL_TERMINAL_PROCESSING:
             return True
-    except OSError:
+    # Can get TypeError in testsuite where 'fd' is a Mock()
+    except (OSError, TypeError):
         return False

demos/demo01.py

@@ -10,9 +10,9 @@
 # Add parent dir to sys path, so the following 'import colorama' always finds
 # the local source in preference to any installed version of colorama.
 import fixpath
-from colorama import init, Fore, Back, Style
+from colorama import just_fix_windows_console, Fore, Back, Style
 
-init()
+just_fix_windows_console()
 
 # Fore, Back and Style are convenience classes for the constant ANSI strings that set
 #     the foreground, background and style. The don't have any magic of their own.

demos/demo02.py

@@ -5,9 +5,9 @@
 
 from __future__ import print_function
 import fixpath
-from colorama import init, Fore, Back, Style
+from colorama import just_fix_windows_console, Fore, Back, Style
 
-init()
+just_fix_windows_console()
 
 print(Fore.GREEN + 'green, '
     + Fore.RED + 'red, '

demos/demo06.py

@@ -24,7 +24,7 @@
 PASSES = 1000
 
 def main():
-    colorama.init()
+    colorama.just_fix_windows_console()
     pos = lambda y, x: Cursor.POS(x, y)
     # draw a white border.
     print(Back.WHITE, end='')

demos/demo07.py

@@ -16,7 +16,7 @@ def main():
     aba
     3a4
     """
-    colorama.init()
+    colorama.just_fix_windows_console()
     print("aaa")
     print("aaa")
     print("aaa")