main: refactor color detection

Signed-off-by: Filipe Laíns <lains@riseup.net>

Author: FFY00

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,8 +33,20 @@
     'underline': '\33[4m',
     'reset': '\33[0m',
 }
-if 'FORCE_COLOR' not in os.environ and not sys.stdout.isatty() or 'NO_COLOR' in os.environ:
-    _STYLES = {color: '' for color in _STYLES}
+_NO_COLORS = {color: '' for color in _COLORS}
+
+
+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(

tests/test_main.py

@@ -338,6 +338,7 @@ def test_output_env_subprocess_error(
     stdout_body,
     stdout_error,
 ):
+    monkeypatch.delenv('NO_COLOR')
     monkeypatch.setenv('FORCE_COLOR' if color else 'NO_COLOR', '')
 
     importlib.reload(build.__main__)  # reload module to set _STYLES
@@ -352,3 +353,34 @@ def test_output_env_subprocess_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