Add --color-output=force that colors output regardless if it is going to a terminal

docs/source/command_line.rst

@@ -923,9 +923,30 @@ in error messages.
     Use visually nicer output in error messages: use soft word wrap,
     show source code snippets, and show error location markers.
 
+.. option:: --color-output[=force]
+
+    ``--color-output`` enables colored output if the output is going to a terminal.
+
+    ``--color-output=force`` enables colored output unconditionally.
+
+    Output will still be uncolored if mypy fails to detect a color code scheme.
+
+    .. note::
+        When the environment variable ``MYPY_FORCE_COLOR`` is set to a
+        non-``0`` non-empty string, mypy ignores ``--color-output[=force]``
+        and ``--no-color-output``, and behaves as if ``--color-output=force``
+        is given.
+
+        If ``MYPY_FORCE_COLOR`` is ``0``, it has no effect.
+
+        If ``MYPY_FORCE_COLOR`` is not defined, but ``FORCE_COLOR`` is defined,
+        it is treated the same way (like a fallback).
+
 .. option:: --no-color-output
 
-    This flag will disable color output in error messages, enabled by default.
+    Disables colored output.
+
+    See also note above.
 
 .. option:: --no-error-summary
 

mypy/dmypy_server.py

@@ -24,6 +24,7 @@
 import mypy.build
 import mypy.errors
 import mypy.main
+from mypy import util
 from mypy.dmypy_util import WriteToConn, receive, send
 from mypy.find_sources import InvalidSourceList, create_source_list
 from mypy.fscache import FileSystemCache
@@ -201,7 +202,13 @@ def __init__(self, options: Options, status_file: str, timeout: int | None = Non
 
         # Since the object is created in the parent process we can check
         # the output terminal options here.
-        self.formatter = FancyFormatter(sys.stdout, sys.stderr, options.hide_error_codes)
+        self.formatter = FancyFormatter(
+            sys.stdout,
+            sys.stderr,
+            options.hide_error_codes,
+            color_request=util.should_force_color() or options.color_output is not False,
+            warn_color_fail=options.warn_color_fail,
+        )
 
     def _response_metadata(self) -> dict[str, str]:
         py_version = f"{self.options.python_version[0]}_{self.options.python_version[1]}"
@@ -841,7 +848,11 @@ def pretty_messages(
         is_tty: bool = False,
         terminal_width: int | None = None,
     ) -> list[str]:
-        use_color = self.options.color_output and is_tty
+        use_color = (
+            True
+            if util.should_force_color() or self.options.color_output == "force"
+            else (is_tty if self.options.color_output is True else False)
+        )
         fit_width = self.options.pretty and is_tty
         if fit_width:
             messages = self.formatter.fit_in_terminal(

mypy/main.py

@@ -91,7 +91,19 @@ def main(
         options.fast_exit = False
 
     formatter = util.FancyFormatter(
-        stdout, stderr, options.hide_error_codes, hide_success=bool(options.output)
+        stdout,
+        stderr,
+        options.hide_error_codes,
+        hide_success=bool(options.output),
+        color_request=util.should_force_color() or options.color_output is not False,
+        warn_color_fail=options.warn_color_fail,
+    )
+
+    # Type annotation needed for mypy (Pyright understands this)
+    use_color: bool = (
+        True
+        if util.should_force_color() or options.color_output == "force"
+        else formatter.default_colored if options.color_output is True else False
     )
 
     if options.allow_redefinition_new and not options.local_partial_types:
@@ -124,7 +136,7 @@ def main(
         install_types(formatter, options, non_interactive=options.non_interactive)
         return
 
-    res, messages, blockers = run_build(sources, options, fscache, t0, stdout, stderr)
+    res, messages, blockers = run_build(sources, options, fscache, t0, stdout, stderr, use_color)
 
     if options.non_interactive:
         missing_pkgs = read_types_packages_to_install(options.cache_dir, after_run=True)
@@ -133,8 +145,10 @@ def main(
             install_types(formatter, options, after_run=True, non_interactive=True)
             fscache.flush()
             print()
-            res, messages, blockers = run_build(sources, options, fscache, t0, stdout, stderr)
-        show_messages(messages, stderr, formatter, options)
+            res, messages, blockers = run_build(
+                sources, options, fscache, t0, stdout, stderr, use_color
+            )
+        show_messages(messages, stderr, formatter, options, use_color)
 
     if MEM_PROFILE:
         from mypy.memprofile import print_memory_profile
@@ -148,12 +162,12 @@ def main(
     if options.error_summary:
         if n_errors:
             summary = formatter.format_error(
-                n_errors, n_files, len(sources), blockers=blockers, use_color=options.color_output
+                n_errors, n_files, len(sources), blockers=blockers, use_color=use_color
             )
             stdout.write(summary + "\n")
         # Only notes should also output success
         elif not messages or n_notes == len(messages):
-            stdout.write(formatter.format_success(len(sources), options.color_output) + "\n")
+            stdout.write(formatter.format_success(len(sources), use_color) + "\n")
         stdout.flush()
 
     if options.install_types and not options.non_interactive:
@@ -182,9 +196,15 @@ def run_build(
     t0: float,
     stdout: TextIO,
     stderr: TextIO,
+    use_color: bool,
 ) -> tuple[build.BuildResult | None, list[str], bool]:
     formatter = util.FancyFormatter(
-        stdout, stderr, options.hide_error_codes, hide_success=bool(options.output)
+        stdout,
+        stderr,
+        options.hide_error_codes,
+        hide_success=bool(options.output),
+        color_request=use_color,
+        warn_color_fail=options.warn_color_fail,
     )
 
     messages = []
@@ -200,7 +220,7 @@ def flush_errors(filename: str | None, new_messages: list[str], serious: bool) -
             # Collect messages and possibly show them later.
             return
         f = stderr if serious else stdout
-        show_messages(new_messages, f, formatter, options)
+        show_messages(new_messages, f, formatter, options, use_color)
 
     serious = False
     blockers = False
@@ -238,10 +258,14 @@ def flush_errors(filename: str | None, new_messages: list[str], serious: bool) -
 
 
 def show_messages(
-    messages: list[str], f: TextIO, formatter: util.FancyFormatter, options: Options
+    messages: list[str],
+    f: TextIO,
+    formatter: util.FancyFormatter,
+    options: Options,
+    use_color: bool,
 ) -> None:
     for msg in messages:
-        if options.color_output:
+        if use_color:
             msg = formatter.colorize(msg)
         f.write(msg + "\n")
     f.flush()
@@ -462,6 +486,19 @@ def __call__(
         parser.exit()
 
 
+# Coupled with the usage in define_options
+class ColorOutputAction(argparse.Action):
+    def __call__(
+        self,
+        parser: argparse.ArgumentParser,
+        namespace: argparse.Namespace,
+        values: str | Sequence[Any] | None,
+        option_string: str | None = None,
+    ) -> None:
+        assert values in ("force", None)
+        setattr(namespace, self.dest, True if values is None else "force")
+
+
 def define_options(
     program: str = "mypy",
     header: str = HEADER,
@@ -993,11 +1030,26 @@ def add_invertible_flag(
         " and show error location markers",
         group=error_group,
     )
-    add_invertible_flag(
-        "--no-color-output",
+    # XXX Setting default doesn't seem to work unless I change
+    # the attribute in options.Options
+    error_group.add_argument(
+        "--color-output",
         dest="color_output",
-        default=True,
-        help="Do not colorize error messages",
+        action=ColorOutputAction,
+        nargs="?",
+        choices=["force"],
+        help="Colorize error messages if output is going to a terminal (inverse: --no-color-output). "
+        "When --color-output=auto, colorizes output unconditionally",
+    )
+    error_group.add_argument(
+        "--no-color-output", dest="color_output", action="store_false", help=argparse.SUPPRESS
+    )
+    add_invertible_flag(
+        "--warn-color-fail",
+        dest="warn_color_fail",
+        default=False,
+        help="Print warning message when mypy cannot detect "
+        "a terminal color scheme and colored output is requested",
         group=error_group,
     )
     add_invertible_flag(
@@ -1530,7 +1582,8 @@ def set_strict_flags() -> None:
                 reason = cache.find_module(p)
                 if reason is ModuleNotFoundReason.FOUND_WITHOUT_TYPE_HINTS:
                     fail(
-                        f"Package '{p}' cannot be type checked due to missing py.typed marker. See https://mypy.readthedocs.io/en/stable/installed_packages.html for more details",
+                        f"Package '{p}' cannot be type checked due to missing py.typed marker. "
+                        "See https://mypy.readthedocs.io/en/stable/installed_packages.html for more details",
                         stderr,
                         options,
                     )

mypy/options.py

@@ -7,7 +7,7 @@
 import warnings
 from collections.abc import Mapping
 from re import Pattern
-from typing import Any, Callable, Final
+from typing import Any, Callable, Final, Literal
 
 from mypy import defaults
 from mypy.errorcodes import ErrorCode, error_codes
@@ -206,8 +206,9 @@ def __init__(self) -> None:
         self.show_error_context = False
 
         # Use nicer output (when possible).
-        self.color_output = True
+        self.color_output: bool | Literal["force"] = True
         self.error_summary = True
+        self.warn_color_fail: bool = False
 
         # Assume arguments with default values of None are Optional
         self.implicit_optional = False

mypy/test/test_color_output.py

@@ -0,0 +1,82 @@
+import subprocess
+import sys
+from typing import TYPE_CHECKING
+
+import pytest
+
+# XXX Would like help with this test, how do I make it runnable?
+# Haven't run this test yet
+
+PTY_SIZE = (80, 40)
+
+if sys.platform == "win32":
+    if TYPE_CHECKING:
+        # This helps my IDE find the type annotations
+        from winpty.winpty import PTY  # type:ignore[import-untyped]
+    else:
+        from winpty import PTY
+
+    def run_pty(cmd: str, env: dict[str, str] = {}) -> tuple[str, str]:
+        pty = PTY(*PTY_SIZE)
+        # For the purposes of this test, str.split() is enough
+        appname, cmdline = cmd.split(maxsplit=1)
+        pty.spawn(appname, cmdline, "\0".join(map(lambda kv: f"{kv[0]}={kv[1]}", env.items())))
+        while pty.isalive():
+            pass
+        return pty.read(), pty.read_stderr()
+
+elif sys.platform == "unix":
+    from pty import openpty
+
+    def run_pty(cmd: str, env: dict[str, str] = {}) -> tuple[str, str]:
+        # NOTE Would like help checking quality of this function,
+        # it's partially written by Copilot because I'm not familiar with Unix openpty
+        # and cannot use Unix openpty
+        master_fd, slave_fd = openpty()
+        try:
+            p = subprocess.run(cmd, stdout=slave_fd, stderr=subprocess.PIPE, env=env, text=True)
+            os.close(slave_fd)
+            return os.read(slave_fd, 10000).decode(), p.stderr
+        finally:
+            os.close(master_fd)
+
+
+def test(expect_color: bool, pty: bool, cmd: str, env: dict[str, str] = {}) -> None:
+    if pty:
+        stdout, stderr = run_pty(cmd, env=env)
+    else:
+        proc = subprocess.run(cmd, capture_output=True, env=env, text=True)
+        stdout = proc.stdout
+        stderr = proc.stderr
+    if "Found" not in stdout:  # ??
+        pytest.fail("Command failed to complete or did not detect type error")
+    if expect_color:  # Expect color control chars
+        assert "<string>:1: error:" not in stdout
+        assert "\nFound" not in stdout
+    else:  # Expect no color control chars
+        assert "<string>:1: error:" in stdout
+        assert "\nFound" in stdout
+
+
+def test_pty(expect_color: bool, cmd: str, env: dict[str, str] = {}) -> None:
+    test(expect_color, True, cmd, env)
+
+
+def test_not_pty(expect_color: bool, cmd: str, env: dict[str, str] = {}) -> None:
+    test(expect_color, False, cmd, env)
+
+
+@pytest.mark.parametrize("command", ["mypy", "dmypy run --"])
+def test_it(command: str) -> None:
+    # Note: Though we don't check stderr, capturing it is useful
+    # because it provides traceback if mypy crashes due to exception
+    # and pytest reveals it upon failure (?)
+    test_pty(True, f"{command} -c \"1+'a'\" --color-output=force")
+    test_pty(False, f"{command} -c \"1+'a'\" --no-color-output")
+    test_not_pty(False, f"{command} -c \"1+'a'\" --color-output")
+    test_not_pty(True, f"{command} -c \"1+'a'\" --color-output=force")
+    test_not_pty(False, f"{command} -c \"1+'a'\" --color-output", {"MYPY_FORCE_COLOR": "1"})
+    test_not_pty(True, f"{command} -c \"1+'a'\" --color-output=force", {"MYPY_FORCE_COLOR": "1"})
+    test_not_pty(False, f"{command} -c \"1+'a'\" --no-color-output", {"MYPY_FORCE_COLOR": "1"})
+    test_not_pty(False, f"{command} -c \"1+'a'\" --no-color-output", {"FORCE_COLOR": "1"})
+    test_not_pty(False, f"{command} -c \"1+'a'\" --color-output", {"MYPY_FORCE_COLOR": "0"})