sts-sponsors team mailing list archive

[Adam Collard <mp+434512@xxxxxxxxxxxxxxxxxx>]

Adam Collard has proposed merging ~adam-collard/maas:3.3-colorama-strikes-back into maas:3.3.

Commit message:
get original stdout from Colorama

When colorama is present, Django loads it automatically. This module
interferes with the CLI raw output, so we need to get the original
stdout descriptor.

Cherry pick of dc4621928f2d3713dd0c632e98d4ecf1491f5b77

Requested reviews:
  Adam Collard (adam-collard)

For more details, see:
https://code.launchpad.net/~adam-collard/maas/+git/maas/+merge/434512
-- 
Your team MAAS Committers is subscribed to branch maas:3.3.

diff --git a/src/maascli/__init__.py b/src/maascli/__init__.py
index fc57914..e309b21 100644
--- a/src/maascli/__init__.py
+++ b/src/maascli/__init__.py
@@ -25,25 +25,12 @@ def snap_setup():
         )
 
 
-def disable_colorama():
-    """Django pulls Colorama automatically if it's in PYTHONPATH, which
-    interferes with our raw output mode"""
-    try:
-        import colorama
-
-        colorama.deinit()
-    except (ImportError, OSError):
-        pass
-
-
 def main(argv=sys.argv):
     # If no arguments have been passed be helpful and point out --help.
     verbose_errors = "MAAS_CLI_VERBOSE_ERRORS" in os.environ
 
     snap_setup()
 
-    disable_colorama()
-
     parser = prepare_parser(argv)
 
     try:
diff --git a/src/maascli/tests/test_integration.py b/src/maascli/tests/test_integration.py
index 9cef775..4f61d1a 100644
--- a/src/maascli/tests/test_integration.py
+++ b/src/maascli/tests/test_integration.py
@@ -8,7 +8,7 @@ import os.path
 import random
 from subprocess import CalledProcessError, check_output, STDOUT
 
-from maascli import disable_colorama, main
+from maascli import main
 from maascli.config import ProfileConfig
 from maascli.testing.config import make_configs
 from maascli.utils import handler_command_name
@@ -74,9 +74,6 @@ class TestMain(MAASTestCase):
         handler_name = handler_command_name(resource_name)
         command = "maas", profile_name, handler_name
 
-        # Trying to disable colorama conflicts with the CaptureStandardIO below
-        self.patch(disable_colorama)
-
         with CaptureStandardIO() as stdio:
             error = self.assertRaises(SystemExit, main, command)
 
diff --git a/src/maascli/utils.py b/src/maascli/utils.py
index f35e929..f38aa5b 100644
--- a/src/maascli/utils.py
+++ b/src/maascli/utils.py
@@ -150,6 +150,22 @@ def print_response_headers(headers, file=None):
         print(form % (cap(header), headers[header]), file=file)
 
 
+def get_orig_stdout(file):
+    # Get the underlying buffer if we're writing to stdout. This allows us to
+    # write bytes directly, without attempting to convert the bytes to unicode.
+    # Unicode output may not be desired; the HTTP response could be raw bytes.
+    try:
+        import colorama
+
+        if isinstance(file, colorama.ansitowin32.StreamWrapper):
+            file = colorama.initialise.orig_stdout
+    except (ImportError, OSError):
+        pass
+    if isinstance(file, io.TextIOWrapper):
+        file = file.buffer
+    return file
+
+
 def print_response_content(response, content, file=None):
     """Write the response's content to stdout.
 
@@ -159,14 +175,10 @@ def print_response_content(response, content, file=None):
     :param file: a binary stream opened for writing (optional)
     """
     file = sys.stdout if file is None else file
+    file = get_orig_stdout(file)
     is_tty = file.isatty()
     success = response.status // 100 == 2
     is_textual = is_response_textual(response)
-    # Get the underlying buffer if we're writing to stdout. This allows us to
-    # write bytes directly, without attempting to convert the bytes to unicode.
-    # Unicode output may not be desired; the HTTP response could be raw bytes.
-    if isinstance(file, io.TextIOWrapper):
-        file = file.buffer
     if is_tty and success and is_textual:
         file.write(b"Success.\n")
         file.write(b"Machine-readable output follows:\n")