Repository: pypa/build
Pull request #335 (2021-07-28 16:41:18)

main: fix replacing color codes in tracebacks
Signed-off-by: Filipe Laíns <lains@riseup.net>

FFY00: I think adding a test for this that would actually be helpful in development would be very tricky, and I feel like it might not be worth the time. The code that we would be covering would only generate formatting errors, not functional ones which would prevent people from working.
How do you test this kind of stuff in tox?
gaborbernat: Why would it be tricky? If now the contract is that the output is nicely colored, why would a test validating that be not worth it? tox threats output similarly. We run the code in a subprocess, and then assert the stdout/stderr contains the right series of bytes. If you can run it without needing a new subprocess you can assert on capsys/capfd and do the same. Doesn't have to test color output for all functional cases, just one run that covers all cases. Bugs related to the coloring helper function can be tested by directly calling the helper function from your test.
layday: It's not entirely clear what this is fixing. Are you trying to avoid replacing "[dim]" and "[reset]" in the traceback itself?
FFY00: Yes, as @henryiii pointed out.
FFY00: The tests are failing on Windows and I don't really get why. Since I don't have any Windows setup right now, it might be a little while before I am able to set one up and debug the issues. If anyone is able to look into it, I would really appreciate it.
@henryiii you use Windows, right? Perhaps you could help out if you have time.
FFY00: It was colorama. It's a bit weird since this test was passing previously.

Patch:
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index d0b38c6..19d4c29 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -8,6 +8,7 @@ Unreleased
 
 - Improved output (`PR #333`_, Fixes `#142`_)
 - The CLI now honnors `NO_COLOR`_ (`PR #333`_)
+- The CLI can now be forced to colorize the output by setting the ``FORCE_COLOR`` environment variable (`PR #335`_)
 - Added logging to ``build`` and ``build.env`` (`PR #333`_)
 
 
@@ -17,6 +18,7 @@ Breaking Changes
 - Dropped support for Python 2 and 3.5.
 
 .. _PR #333: https://github.com/pypa/build/pull/333
+.. _PR #335: https://github.com/pypa/build/pull/335
 .. _#142: https://github.com/pypa/build/issues/142
 .. _NO_COLOR: https://no-color.org
 
diff --git a/src/build/__main__.py b/src/build/__main__.py
index 6fc6660..8adc8ac 100644
--- a/src/build/__main__.py
+++ b/src/build/__main__.py
@@ -13,7 +13,7 @@ import textwrap
 import traceback
 import warnings
 
-from typing import Iterable, Iterator, List, Optional, Sequence, TextIO, Type, Union
+from typing import Dict, Iterable, Iterator, List, Optional, Sequence, TextIO, Type, Union
 
 import build
 
@@ -24,7 +24,7 @@ from build.env import IsolatedEnvBuilder
 __all__ = ['build', 'main', 'main_parser']
 
 
-_STYLES = {
+_COLORS = {
     'red': '\33[91m',
     'green': '\33[92m',
     'yellow': '\33[93m',
@@ -33,13 +33,20 @@ _STYLES = {
     'underline': '\33[4m',
     'reset': '\33[0m',
 }
+_NO_COLORS = {color: '' for color in _COLORS}
 
 
-def _print(message: str) -> None:
-    color = sys.stdout.isatty() and 'NO_COLOR' not in os.environ
-    for name, code in _STYLES.items():
-        message = message.replace(f'[{name}]', code if color else '')
-    print(message)
+def _init_colors() -> Dict[str, str]:
+    if 'NO_COLOR' in os.environ:
+        if 'FORCE_COLOR' in os.environ:
+            warnings.warn('Both NO_COLOR and FORCE_COLOR environment variables are set, disabling color')
+        return _NO_COLORS
+    elif 'FORCE_COLOR' in os.environ or sys.stdout.isatty():
+        return _COLORS
+    return _NO_COLORS
+
+
+_STYLES = _init_colors()
 
 
 def _showwarning(
@@ -50,7 +57,7 @@ def _showwarning(
     file: Optional[TextIO] = None,
     line: Optional[str] = None,
 ) -> None:  # pragma: no cover
-    _print(f'[yellow]WARNING[reset] {message}')
+    print('{yellow}WARNING{reset} {}'.format(message, **_STYLES))
 
 
 def _setup_cli() -> None:
@@ -71,20 +78,20 @@ def _error(msg: str, code: int = 1) -> None:  # pragma: no cover
     :param msg: Error message
     :param code: Error code
     """
-    _print(f'[red]ERROR[reset] {msg}')
+    print('{red}ERROR{reset} {}'.format(msg, **_STYLES))
     exit(code)
 
 
 class _ProjectBuilder(ProjectBuilder):
     @staticmethod
     def log(message: str) -> None:
-        _print(f'[bold]* {message}[reset]')
+        print('{bold}* {}{reset}'.format(message, **_STYLES))
 
 
 class _IsolatedEnvBuilder(IsolatedEnvBuilder):
     @staticmethod
     def log(message: str) -> None:
-        _print(f'[bold]* {message}[reset]')
+        print('{bold}* {}{reset}'.format(message, **_STYLES))
 
 
 def _format_dep_chain(dep_chain: Sequence[str]) -> str:
@@ -115,7 +122,7 @@ def _build_in_current_env(
         missing = builder.check_dependencies(distribution)
         if missing:
             dependencies = ''.join('\n\t' + dep for deps in missing for dep in (deps[0], _format_dep_chain(deps[1:])) if dep)
-            _error(f'Missing dependencies:{dependencies}')
+            _error(f'\nMissing dependencies:{dependencies}')
 
     return builder.build(distribution, outdir, config_settings or {})
 
@@ -154,7 +161,7 @@ def _handle_build_error() -> Iterator[None]:
                 tb = ''.join(tb_lines)
             else:
                 tb = traceback.format_exc(-1)
-            _print('\n[dim]{}[reset]\n'.format(tb.strip('\n')))
+            print('\n{dim}{}{reset}\n'.format(tb.strip('\n'), **_STYLES))
         _error(str(e))
 
 
@@ -364,10 +371,13 @@ def main(cli_args: Sequence[str], prog: Optional[str] = None) -> None:  # noqa:
             built = build_call(
                 args.srcdir, outdir, distributions, config_settings, not args.no_isolation, args.skip_dependency_check
             )
-            artifact_list = _natural_language_list([f'[underline]{artifact}[reset][bold][green]' for artifact in built])
-            _print(f'[bold][green]Successfully built {artifact_list}')
+            artifact_list = _natural_language_list(
+                ['{underline}{}{reset}{bold}{green}'.format(artifact, **_STYLES) for artifact in built]
+            )
+            print('{bold}{green}Successfully built {}{reset}'.format(artifact_list, **_STYLES))
     except Exception as e:  # pragma: no cover
-        _print('\n[dim]{}[reset]\n'.format(traceback.format_exc().strip('\n')))
+        tb = traceback.format_exc().strip('\n')
+        print('\n{dim}{}{reset}\n'.format(tb, **_STYLES))
         _error(str(e))
 
 


Test patch:
diff --git a/tests/test_main.py b/tests/test_main.py
index 253c30d..86b0176 100644
--- a/tests/test_main.py
+++ b/tests/test_main.py
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: MIT
 
 import contextlib
+import importlib
 import io
 import os
 import re
@@ -161,7 +162,7 @@ def test_build_no_isolation_with_check_deps(mocker, test_flit_path, missing_deps
     build.__main__.build_package(test_flit_path, '.', ['sdist'], isolation=False)
 
     build_cmd.assert_called_with('sdist', '.', {})
-    error.assert_called_with('Missing dependencies:' + output)
+    error.assert_called_with('\nMissing dependencies:' + output)
 
 
 @pytest.mark.isolated
@@ -291,19 +292,101 @@ def test_output(test_setuptools_path, tmp_dir, capsys, args, output):
     assert stdout.splitlines() == output
 
 
-def test_output_env_subprocess_error(test_invalid_requirements_path, tmp_dir, capsys):
+@pytest.fixture()
+def main_reload_styles():
+    try:
+        yield
+    finally:
+        importlib.reload(build.__main__)
+
+
+@pytest.mark.parametrize(
+    ('color', 'stdout_error', 'stdout_body'),
+    [
+        (
+            False,
+            'ERROR ',
+            [
+                '* Creating venv isolated environment...',
+                '* Installing packages in isolated environment... (setuptools >= 42.0.0, this is invalid, wheel >= 0.36.0)',
+                '',
+                'Traceback (most recent call last):',
+            ],
+        ),
+        (
+            True,
+            '\33[91mERROR\33[0m ',
+            [
+                '\33[1m* Creating venv isolated environment...\33[0m',
+                '\33[1m* Installing packages in isolated environment... '
+                '(setuptools >= 42.0.0, this is invalid, wheel >= 0.36.0)\33[0m',
+                '',
+                '\33[2mTraceback (most recent call last):',
+            ],
+        ),
+    ],
+    ids=['no-color', 'color'],
+)
+def test_output_env_subprocess_error(
+    mocker,
+    monkeypatch,
+    main_reload_styles,
+    test_invalid_requirements_path,
+    tmp_dir,
+    capsys,
+    color,
+    stdout_body,
+    stdout_error,
+):
+    try:
+        # do not inject hook to have clear output on capsys
+        mocker.patch('colorama.init')
+    except ModuleNotFoundError:  # colorama might not be available
+        pass
+
+    monkeypatch.delenv('NO_COLOR', raising=False)
+    monkeypatch.setenv('FORCE_COLOR' if color else 'NO_COLOR', '')
+
+    importlib.reload(build.__main__)  # reload module to set _STYLES
+
     with pytest.raises(SystemExit):
         build.__main__.main([test_invalid_requirements_path, '-o', tmp_dir])
     stdout, stderr = capsys.readouterr()
     stdout, stderr = stdout.splitlines(), stderr.splitlines()
 
-    assert stdout[:4] == [
-        '* Creating venv isolated environment...',
-        '* Installing packages in isolated environment... (setuptools >= 42.0.0, this is invalid, wheel >= 0.36.0)',
-        '',
-        'Traceback (most recent call last):',
-    ]
-    assert stdout[-1].startswith('ERROR ')
+    assert stdout[:4] == stdout_body
+    assert stdout[-1].startswith(stdout_error)
 
     assert len(stderr) == 1
     assert stderr[0].startswith('ERROR: Invalid requirement: ')
+
+
+@pytest.mark.parametrize(
+    ('tty', 'env', 'colors'),
+    [
+        (False, {}, build.__main__._NO_COLORS),
+        (True, {}, build.__main__._COLORS),
+        (False, {'NO_COLOR': ''}, build.__main__._NO_COLORS),
+        (True, {'NO_COLOR': ''}, build.__main__._NO_COLORS),
+        (False, {'FORCE_COLOR': ''}, build.__main__._COLORS),
+        (True, {'FORCE_COLOR': ''}, build.__main__._COLORS),
+    ],
+)
+def test_colors(mocker, monkeypatch, main_reload_styles, tty, env, colors):
+    mocker.patch('sys.stdout.isatty', return_value=tty)
+    for key, value in env.items():
+        monkeypatch.setenv(key, value)
+
+    importlib.reload(build.__main__)  # reload module to set _STYLES
+
+    assert build.__main__._STYLES == colors
+
+
+def test_colors_conflict(monkeypatch, main_reload_styles):
+    monkeypatch.setenv('NO_COLOR', '')
+    monkeypatch.setenv('FORCE_COLOR', '')
+
+    with pytest.warns(Warning, match='Both NO_COLOR and FORCE_COLOR environment variables are set, disabling color'):
+        importlib.reload(build.__main__)
+
+    assert build.__main__._STYLES == build.__main__._NO_COLORS