main: fix replacing color codes in tracebacks

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
 

src/build/__main__.py

@@ -13,7 +13,7 @@
 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 @@
 __all__ = ['build', 'main', 'main_parser']
 
 
-_STYLES = {
+_COLORS = {
     'red': '\33[91m',
     'green': '\33[92m',
     'yellow': '\33[93m',
@@ -33,13 +33,20 @@
     '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))
 
 

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