Keep styles in context instead of reloading the module

src/build/__main__.py

@@ -4,6 +4,7 @@
 
 import argparse
 import contextlib
+import contextvars
 import os
 import platform
 import shutil
@@ -38,21 +39,21 @@
 _NO_COLORS = {color: '' for color in _COLORS}
 
 
-def _init_colors() -> dict[str, str]:
+_styles = contextvars.ContextVar('_styles', default=_COLORS)
+
+
+def _init_colors() -> None:
     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', stacklevel=2)
-        return _NO_COLORS
+        _styles.set(_NO_COLORS)
     elif 'FORCE_COLOR' in os.environ or sys.stdout.isatty():
-        return _COLORS
-    return _NO_COLORS
-
-
-_STYLES = _init_colors()
+        return
+    _styles.set(_NO_COLORS)
 
 
 def _cprint(fmt: str = '', msg: str = '', file: TextIO | None = None) -> None:
-    print(fmt.format(msg, **_STYLES), file=file, flush=True)
+    print(fmt.format(msg, **_styles.get()), file=file, flush=True)
 
 
 def _showwarning(
@@ -98,6 +99,8 @@ def _setup_cli(*, verbosity: int) -> None:
         except ModuleNotFoundError:
             pass
 
+    _init_colors()
+
     _ctx.LOGGER.set(_log)
     _ctx.VERBOSITY.set(verbosity)
 
@@ -415,7 +418,7 @@ def main(cli_args: Sequence[str], prog: str | None = None) -> None:
             args.srcdir, outdir, distributions, config_settings, not args.no_isolation, args.skip_dependency_check
         )
         artifact_list = _natural_language_list(
-            ['{underline}{}{reset}{bold}{green}'.format(artifact, **_STYLES) for artifact in built]
+            ['{underline}{}{reset}{bold}{green}'.format(artifact, **_styles.get()) for artifact in built]
         )
         _cprint('{bold}{green}Successfully built {}{reset}', artifact_list)
 

tests/test_main.py

@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: MIT
 
 import contextlib
-import importlib
 import io
 import os
 import re
@@ -312,14 +311,6 @@ def test_output(package_test_setuptools, tmp_dir, capsys, args, output):
     assert stdout.splitlines() == output
 
 
-@pytest.fixture()
-def main_reload_styles():
-    try:
-        yield
-    finally:
-        importlib.reload(build.__main__)
-
-
 @pytest.mark.pypy3323bug
 @pytest.mark.parametrize(
     ('color', 'stdout_error', 'stdout_body'),
@@ -351,7 +342,6 @@ def main_reload_styles():
 def test_output_env_subprocess_error(
     mocker,
     monkeypatch,
-    main_reload_styles,
     package_test_invalid_requirements,
     tmp_dir,
     capsys,
@@ -368,8 +358,6 @@ def test_output_env_subprocess_error(
     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([package_test_invalid_requirements, '-o', tmp_dir])
     stdout, stderr = capsys.readouterr()
@@ -394,17 +382,17 @@ def test_output_env_subprocess_error(
         (True, {'FORCE_COLOR': ''}, build.__main__._COLORS),
     ],
 )
-def test_colors(mocker, monkeypatch, main_reload_styles, tty, env, colors):
+def test_colors(mocker, monkeypatch, 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
+    build.__main__._init_colors()
 
-    assert build.__main__._STYLES == colors
+    assert build.__main__._styles.get() == colors
 
 
-def test_colors_conflict(monkeypatch, main_reload_styles):
+def test_colors_conflict(monkeypatch):
     with monkeypatch.context() as m:
         m.setenv('NO_COLOR', '')
         m.setenv('FORCE_COLOR', '')
@@ -413,9 +401,9 @@ def test_colors_conflict(monkeypatch, main_reload_styles):
             UserWarning,
             match='Both NO_COLOR and FORCE_COLOR environment variables are set, disabling color',
         ):
-            importlib.reload(build.__main__)
+            build.__main__._init_colors()
 
-        assert build.__main__._STYLES == build.__main__._NO_COLORS
+        assert build.__main__._styles.get() == build.__main__._NO_COLORS
 
 
 def test_venv_fail(monkeypatch, package_test_flit, tmp_dir, capsys):
@@ -425,8 +413,6 @@ def raise_called_process_err(*args, **kwargs):
     monkeypatch.setattr(venv.EnvBuilder, 'create', raise_called_process_err)
     monkeypatch.setenv('NO_COLOR', '')
 
-    importlib.reload(build.__main__)  # reload module to set _STYLES
-
     with pytest.raises(SystemExit):
         build.__main__.main([package_test_flit, '-o', tmp_dir])