main: fix replacing color codes in tracebacks

src/build/__main__.py

@@ -33,13 +33,8 @@
     'underline': '\33[4m',
     'reset': '\33[0m',
 }
-
-
-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)
+if not sys.stdout.isatty() or 'NO_COLOR' in os.environ:
+    _STYLES = {color: '' for color in _STYLES}
 
 
 def _showwarning(
@@ -50,7 +45,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 +66,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:
@@ -154,7 +149,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 +359,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
@@ -291,19 +292,65 @@ 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,
+):
+    mocker.patch('sys.stdout.isatty', return_value=True)
+    if not color:
+        monkeypatch.setenv('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: ')