Improve control over ANSI output (#6858)

* Move global console_handler into function scope

Signed-off-by: Mike Seplowitz <mseplowitz@bloomberg.net>

* Improve control over ANSI output

- Disabled parallel logger ANSI output if not attached to a tty.
  The console handler and progress stream already checked whether the
  output stream is a tty, but ParallelStreamWriter did not.

- Added --ansi=(never|always|auto) option to allow clearer control over
  ANSI output. Since --no-ansi is the same as --ansi=never, --no-ansi is
  now deprecated.

Signed-off-by: Mike Seplowitz <mseplowitz@bloomberg.net>

compose/cli/colors.py

@@ -1,3 +1,6 @@
+import enum
+import os
+
 from ..const import IS_WINDOWS_PLATFORM
 
 NAMES = [
@@ -12,6 +15,21 @@ NAMES = [
 ]
 
 
+@enum.unique
+class AnsiMode(enum.Enum):
+    """Enumeration for when to output ANSI colors."""
+    NEVER = "never"
+    ALWAYS = "always"
+    AUTO = "auto"
+
+    def use_ansi_codes(self, stream):
+        if self is AnsiMode.ALWAYS:
+            return True
+        if self is AnsiMode.NEVER or os.environ.get('CLICOLOR') == '0':
+            return False
+        return stream.isatty()
+
+
 def get_pairs():
     for i, name in enumerate(NAMES):
         yield (name, str(30 + i))

compose/cli/main.py

@@ -2,7 +2,6 @@ import contextlib
 import functools
 import json
 import logging
-import os
 import pipes
 import re
 import subprocess
@@ -27,6 +26,7 @@ from ..config.types import VolumeSpec
 from ..const import IS_WINDOWS_PLATFORM
 from ..errors import StreamParseError
 from ..metrics.decorator import metrics
+from ..parallel import ParallelStreamWriter
 from ..progress_stream import StreamOutputError
 from ..project import get_image_digests
 from ..project import MissingDigests
@@ -40,6 +40,7 @@ from ..service import ImageType
 from ..service import NeedsBuildError
 from ..service import OperationFailedError
 from ..utils import filter_attached_for_up
+from .colors import AnsiMode
 from .command import get_config_from_options
 from .command import get_project_dir
 from .command import project_from_options
@@ -62,7 +63,6 @@ if not IS_WINDOWS_PLATFORM:
     from dockerpty.pty import PseudoTerminal, RunOperation, ExecOperation
 
 log = logging.getLogger(__name__)
-console_handler = logging.StreamHandler(sys.stderr)
 
 
 def main():  # noqa: C901
@@ -139,18 +139,38 @@ def exit_with_metrics(command, log_msg=None, status=Status.SUCCESS, exit_code=1)
 
 
 def dispatch():
-    setup_logging()
+    console_stream = sys.stderr
+    console_handler = logging.StreamHandler(console_stream)
+    setup_logging(console_handler)
     dispatcher = DocoptDispatcher(
         TopLevelCommand,
         {'options_first': True, 'version': get_version_info('compose')})
 
     options, handler, command_options = dispatcher.parse(sys.argv[1:])
+
+    ansi_mode = AnsiMode.AUTO
+    try:
+        if options.get("--ansi"):
+            ansi_mode = AnsiMode(options.get("--ansi"))
+    except ValueError:
+        raise UserError(
+            'Invalid value for --ansi: {}. Expected one of {}.'.format(
+                options.get("--ansi"),
+                ', '.join(m.value for m in AnsiMode)
+            )
+        )
+    if options.get("--no-ansi"):
+        if options.get("--ansi"):
+            raise UserError("--no-ansi and --ansi cannot be combined.")
+        log.warning('--no-ansi option is deprecated and will be removed in future versions.')
+        ansi_mode = AnsiMode.NEVER
+
     setup_console_handler(console_handler,
                           options.get('--verbose'),
-                          set_no_color_if_clicolor(options.get('--no-ansi')),
+                          ansi_mode.use_ansi_codes(console_handler.stream),
                           options.get("--log-level"))
-    setup_parallel_logger(set_no_color_if_clicolor(options.get('--no-ansi')))
-    if options.get('--no-ansi'):
+    setup_parallel_logger(ansi_mode)
+    if ansi_mode is AnsiMode.NEVER:
         command_options['--no-color'] = True
     return functools.partial(perform_command, options, handler, command_options)
 
@@ -172,7 +192,7 @@ def perform_command(options, handler, command_options):
         handler(command, command_options)
 
 
-def setup_logging():
+def setup_logging(console_handler):
     root_logger = logging.getLogger()
     root_logger.addHandler(console_handler)
     root_logger.setLevel(logging.DEBUG)
@@ -183,14 +203,12 @@ def setup_logging():
     logging.getLogger("docker").propagate = False
 
 
-def setup_parallel_logger(noansi):
-    if noansi:
-        import compose.parallel
-        compose.parallel.ParallelStreamWriter.set_noansi()
+def setup_parallel_logger(ansi_mode):
+    ParallelStreamWriter.set_default_ansi_mode(ansi_mode)
 
 
-def setup_console_handler(handler, verbose, noansi=False, level=None):
-    if handler.stream.isatty() and noansi is False:
+def setup_console_handler(handler, verbose, use_console_formatter=True, level=None):
+    if use_console_formatter:
         format_class = ConsoleWarningFormatter
     else:
         format_class = logging.Formatter
@@ -242,7 +260,8 @@ class TopLevelCommand:
       -c, --context NAME          Specify a context name
       --verbose                   Show more output
       --log-level LEVEL           Set log level (DEBUG, INFO, WARNING, ERROR, CRITICAL)
-      --no-ansi                   Do not print ANSI control characters
+      --ansi (never|always|auto)  Control when to print ANSI control characters
+      --no-ansi                   Do not print ANSI control characters (DEPRECATED)
       -v, --version               Print version and exit
       -H, --host HOST             Daemon socket to connect to
 
@@ -691,7 +710,7 @@ class TopLevelCommand:
         log_printer_from_project(
             self.project,
             containers,
-            set_no_color_if_clicolor(options['--no-color']),
+            options['--no-color'],
             log_args,
             event_stream=self.project.events(service_names=options['SERVICE']),
             keep_prefix=not options['--no-log-prefix']).run()
@@ -1167,7 +1186,7 @@ class TopLevelCommand:
             log_printer = log_printer_from_project(
                 self.project,
                 attached_containers,
-                set_no_color_if_clicolor(options['--no-color']),
+                options['--no-color'],
                 {'follow': True},
                 cascade_stop,
                 event_stream=self.project.events(service_names=service_names),
@@ -1651,7 +1670,3 @@ def warn_for_swarm_mode(client):
             "To deploy your application across the swarm, "
             "use `docker stack deploy`.\n"
         )
-
-
-def set_no_color_if_clicolor(no_color_flag):
-    return no_color_flag or os.environ.get('CLICOLOR') == "0"

compose/parallel.py

@@ -11,6 +11,7 @@ from threading import Thread
 from docker.errors import APIError
 from docker.errors import ImageNotFound
 
+from compose.cli.colors import AnsiMode
 from compose.cli.colors import green
 from compose.cli.colors import red
 from compose.cli.signals import ShutdownException
@@ -83,10 +84,7 @@ def parallel_execute(objects, func, get_name, msg, get_deps=None, limit=None, fa
     objects = list(objects)
     stream = sys.stderr
 
-    if ParallelStreamWriter.instance:
-        writer = ParallelStreamWriter.instance
-    else:
-        writer = ParallelStreamWriter(stream)
+    writer = ParallelStreamWriter.get_or_assign_instance(ParallelStreamWriter(stream))
 
     for obj in objects:
         writer.add_object(msg, get_name(obj))
@@ -259,19 +257,37 @@ class ParallelStreamWriter:
     to jump to the correct line, and write over the line.
     """
 
-    noansi = False
-    lock = Lock()
+    default_ansi_mode = AnsiMode.AUTO
+    write_lock = Lock()
+
     instance = None
+    instance_lock = Lock()
 
     @classmethod
-    def set_noansi(cls, value=True):
-        cls.noansi = value
+    def get_instance(cls):
+        return cls.instance
 
-    def __init__(self, stream):
+    @classmethod
+    def get_or_assign_instance(cls, writer):
+        cls.instance_lock.acquire()
+        try:
+            if cls.instance is None:
+                cls.instance = writer
+            return cls.instance
+        finally:
+            cls.instance_lock.release()
+
+    @classmethod
+    def set_default_ansi_mode(cls, ansi_mode):
+        cls.default_ansi_mode = ansi_mode
+
+    def __init__(self, stream, ansi_mode=None):
+        if ansi_mode is None:
+            ansi_mode = self.default_ansi_mode
         self.stream = stream
+        self.use_ansi_codes = ansi_mode.use_ansi_codes(stream)
         self.lines = []
         self.width = 0
-        ParallelStreamWriter.instance = self
 
     def add_object(self, msg, obj_index):
         if msg is None:
@@ -285,7 +301,7 @@ class ParallelStreamWriter:
         return self._write_noansi(msg, obj_index, '')
 
     def _write_ansi(self, msg, obj_index, status):
-        self.lock.acquire()
+        self.write_lock.acquire()
         position = self.lines.index(msg + obj_index)
         diff = len(self.lines) - position
         # move up
@@ -297,7 +313,7 @@ class ParallelStreamWriter:
         # move back down
         self.stream.write("%c[%dB" % (27, diff))
         self.stream.flush()
-        self.lock.release()
+        self.write_lock.release()
 
     def _write_noansi(self, msg, obj_index, status):
         self.stream.write(
@@ -310,17 +326,10 @@ class ParallelStreamWriter:
     def write(self, msg, obj_index, status, color_func):
         if msg is None:
             return
-        if self.noansi:
-            self._write_noansi(msg, obj_index, status)
-        else:
+        if self.use_ansi_codes:
             self._write_ansi(msg, obj_index, color_func(status))
-
-
-def get_stream_writer():
-    instance = ParallelStreamWriter.instance
-    if instance is None:
-        raise RuntimeError('ParallelStreamWriter has not yet been instantiated')
-    return instance
+        else:
+            self._write_noansi(msg, obj_index, status)
 
 
 def parallel_operation(containers, operation, options, message):

compose/project.py

@@ -789,7 +789,9 @@ class Project:
                 return
 
             try:
-                writer = parallel.get_stream_writer()
+                writer = parallel.ParallelStreamWriter.get_instance()
+                if writer is None:
+                    raise RuntimeError('ParallelStreamWriter has not yet been instantiated')
                 for event in strm:
                     if 'status' not in event:
                         continue

contrib/completion/bash/docker-compose

@@ -164,6 +164,10 @@ _docker_compose_docker_compose() {
 			_filedir "y?(a)ml"
 			return
 			;;
+		--ansi)
+			COMPREPLY=( $( compgen -W "never always auto" -- "$cur" ) )
+			return
+			;;
 		--log-level)
 			COMPREPLY=( $( compgen -W "debug info warning error critical" -- "$cur" ) )
 			return
@@ -616,6 +620,7 @@ _docker_compose() {
 
 	# These options are require special treatment when searching the command.
 	local top_level_options_with_args="
+		--ansi
 		--log-level
 	"
 

contrib/completion/fish/docker-compose.fish

@@ -21,5 +21,7 @@ complete -c docker-compose -l tlscert -r                  -d 'Path to TLS certif
 complete -c docker-compose -l tlskey -r                   -d 'Path to TLS key file'
 complete -c docker-compose -l tlsverify                   -d 'Use TLS and verify the remote'
 complete -c docker-compose -l skip-hostname-check         -d "Don't check the daemon's hostname against the name specified in the client certificate (for example if your docker host is an IP address)"
+complete -c docker-compose -l no-ansi                     -d 'Do not print ANSI control characters'
+complete -c docker-compose -l ansi -a never always auto   -d 'Control when to print ANSI control characters'
 complete -c docker-compose -s h -l help                   -d 'Print usage'
 complete -c docker-compose -s v -l version                -d 'Print version and exit'

contrib/completion/zsh/_docker-compose

@@ -342,6 +342,7 @@ _docker-compose() {
         '--verbose[Show more output]' \
         '--log-level=[Set log level]:level:(DEBUG INFO WARNING ERROR CRITICAL)' \
         '--no-ansi[Do not print ANSI control characters]' \
+        '--ansi=[Control when to print ANSI control characters]:when:(never always auto)' \
         '(-H --host)'{-H,--host}'[Daemon socket to connect to]:host:' \
         '--tls[Use TLS; implied by --tlsverify]' \
         '--tlscacert=[Trust certs signed only by this CA]:ca path:' \

tests/unit/cli/colors_test.py

@@ -0,0 +1,56 @@
+import os
+
+import pytest
+
+from compose.cli.colors import AnsiMode
+from tests import mock
+
+
+@pytest.fixture
+def tty_stream():
+    stream = mock.Mock()
+    stream.isatty.return_value = True
+    return stream
+
+
+@pytest.fixture
+def non_tty_stream():
+    stream = mock.Mock()
+    stream.isatty.return_value = False
+    return stream
+
+
+class TestAnsiModeTestCase:
+
+    @mock.patch.dict(os.environ)
+    def test_ansi_mode_never(self, tty_stream, non_tty_stream):
+        if "CLICOLOR" in os.environ:
+            del os.environ["CLICOLOR"]
+        assert not AnsiMode.NEVER.use_ansi_codes(tty_stream)
+        assert not AnsiMode.NEVER.use_ansi_codes(non_tty_stream)
+
+        os.environ["CLICOLOR"] = "0"
+        assert not AnsiMode.NEVER.use_ansi_codes(tty_stream)
+        assert not AnsiMode.NEVER.use_ansi_codes(non_tty_stream)
+
+    @mock.patch.dict(os.environ)
+    def test_ansi_mode_always(self, tty_stream, non_tty_stream):
+        if "CLICOLOR" in os.environ:
+            del os.environ["CLICOLOR"]
+        assert AnsiMode.ALWAYS.use_ansi_codes(tty_stream)
+        assert AnsiMode.ALWAYS.use_ansi_codes(non_tty_stream)
+
+        os.environ["CLICOLOR"] = "0"
+        assert AnsiMode.ALWAYS.use_ansi_codes(tty_stream)
+        assert AnsiMode.ALWAYS.use_ansi_codes(non_tty_stream)
+
+    @mock.patch.dict(os.environ)
+    def test_ansi_mode_auto(self, tty_stream, non_tty_stream):
+        if "CLICOLOR" in os.environ:
+            del os.environ["CLICOLOR"]
+        assert AnsiMode.AUTO.use_ansi_codes(tty_stream)
+        assert not AnsiMode.AUTO.use_ansi_codes(non_tty_stream)
+
+        os.environ["CLICOLOR"] = "0"
+        assert not AnsiMode.AUTO.use_ansi_codes(tty_stream)
+        assert not AnsiMode.AUTO.use_ansi_codes(non_tty_stream)

tests/unit/cli/main_test.py

@@ -137,21 +137,20 @@ class TestCLIMainTestCase:
 
 class TestSetupConsoleHandlerTestCase:
 
-    def test_with_tty_verbose(self, logging_handler):
+    def test_with_console_formatter_verbose(self, logging_handler):
         setup_console_handler(logging_handler, True)
         assert type(logging_handler.formatter) == ConsoleWarningFormatter
         assert '%(name)s' in logging_handler.formatter._fmt
         assert '%(funcName)s' in logging_handler.formatter._fmt
 
-    def test_with_tty_not_verbose(self, logging_handler):
+    def test_with_console_formatter_not_verbose(self, logging_handler):
         setup_console_handler(logging_handler, False)
         assert type(logging_handler.formatter) == ConsoleWarningFormatter
         assert '%(name)s' not in logging_handler.formatter._fmt
         assert '%(funcName)s' not in logging_handler.formatter._fmt
 
-    def test_with_not_a_tty(self, logging_handler):
-        logging_handler.stream.isatty.return_value = False
-        setup_console_handler(logging_handler, False)
+    def test_without_console_formatter(self, logging_handler):
+        setup_console_handler(logging_handler, False, use_console_formatter=False)
         assert type(logging_handler.formatter) == logging.Formatter
 
 

tests/unit/parallel_test.py

@@ -3,6 +3,7 @@ from threading import Lock
 
 from docker.errors import APIError
 
+from compose.cli.colors import AnsiMode
 from compose.parallel import GlobalLimit
 from compose.parallel import parallel_execute
 from compose.parallel import parallel_execute_iter
@@ -156,7 +157,7 @@ def test_parallel_execute_alignment(capsys):
 
 def test_parallel_execute_ansi(capsys):
     ParallelStreamWriter.instance = None
-    ParallelStreamWriter.set_noansi(value=False)
+    ParallelStreamWriter.set_default_ansi_mode(AnsiMode.ALWAYS)
     results, errors = parallel_execute(
         objects=["something", "something more"],
         func=lambda x: x,
@@ -172,7 +173,7 @@ def test_parallel_execute_ansi(capsys):
 
 def test_parallel_execute_noansi(capsys):
     ParallelStreamWriter.instance = None
-    ParallelStreamWriter.set_noansi()
+    ParallelStreamWriter.set_default_ansi_mode(AnsiMode.NEVER)
     results, errors = parallel_execute(
         objects=["something", "something more"],
         func=lambda x: x,