Repository: hikari-py/hikari
Pull request #682 (2021-06-26 19:54:24)

Default to ANSI colour for Windows Terminal + logic refactor
### Summary
Default to ANSI colour for Windows Terminal + logic refactor
* better fit clicolor's standard behaviour by interpreting os.getenv("CLICOLOR") == "0" as an instruction that colours shouldn't be used rather than passing over it as undefined
* Switch over to patching os.environ in test cases rather than relying on the call order of os.getenv

### Checklist
<!-- Make sure to tick all the following boxes by putting an `x` in between (like this `[x]`) -->
- [x] I have run `nox` and all the pipelines have passed.
- [x] I have made unittests according to the code I have added/modified/deleted.

### Related issues
<!--
To mention an issue use `#issue-id` and to mention a merge request use `!merge-request-id`
To close/fix an issue use `Close #issue-id` or `Fix #issue-id` (depending on the merge request)
-->



Patch:
diff --git a/hikari/internal/ux.py b/hikari/internal/ux.py
index c14aa81d..1c584cae 100644
--- a/hikari/internal/ux.py
+++ b/hikari/internal/ux.py
@@ -136,6 +136,10 @@ def init_logging(
     logging.captureWarnings(True)
 
 
+_UNCONDITIONAL_ANSI_FLAGS: typing.Final[typing.FrozenSet[str]] = frozenset(("PYCHARM_HOSTED", "WT_SESSION"))
+"""Set of env variables which always indicate that ANSI flags should be included."""
+
+
 def print_banner(package: typing.Optional[str], allow_color: bool, force_color: bool) -> None:
     """Print a banner of choice to `sys.stdout`.
 
@@ -223,30 +227,34 @@ def supports_color(allow_color: bool, force_color: bool) -> bool:
     # isatty is not always implemented, https://code.djangoproject.com/ticket/6223
     is_a_tty = hasattr(sys.stdout, "isatty") and sys.stdout.isatty()
 
-    if os.getenv("CLICOLOR_FORCE", "0") != "0" or force_color:
+    clicolor = os.environ.get("CLICOLOR")
+    if os.environ.get("CLICOLOR_FORCE", "0") != "0" or force_color:
         # https://bixense.com/clicolors/
         return True
-    elif os.getenv("CLICOLOR", "0") != "0" and is_a_tty:
+    if clicolor is not None and clicolor != "0" and is_a_tty:
         # https://bixense.com/clicolors/
         return True
-    elif os.getenv("COLORTERM", "").casefold() in ("truecolor", "24bit"):
+    if clicolor == "0":
+        # https://bixense.com/clicolors/
+        return False
+    if os.environ.get("COLORTERM", "").casefold() in ("truecolor", "24bit"):
         # Seems to be used by Gnome Terminal, and Tmpod will beat me if I don't add it.
         # https://gist.github.com/XVilka/8346728#true-color-detection
         return True
 
     plat = sys.platform
-    color_support = False
-
-    if plat != "Pocket PC":
-        if plat == "win32":
-            color_support |= os.getenv("TERM_PROGRAM", None) in ("mintty", "Terminus")
-            color_support |= "ANSICON" in os.environ
-            color_support &= is_a_tty
-        else:
-            color_support = is_a_tty
+    if plat == "Pocket PC":
+        return False
 
-        color_support |= bool(os.getenv("PYCHARM_HOSTED", ""))
+    color_support = False
+    if plat == "win32":
+        color_support |= os.environ.get("TERM_PROGRAM") in ("mintty", "Terminus")
+        color_support |= "ANSICON" in os.environ
+        color_support &= is_a_tty
+    else:
+        color_support = is_a_tty
 
+    color_support |= bool(os.environ.keys() & _UNCONDITIONAL_ANSI_FLAGS)
     return color_support
 
 


Test patch:
diff --git a/tests/hikari/internal/test_ux.py b/tests/hikari/internal/test_ux.py
index 45672c5d..7c76df4b 100644
--- a/tests/hikari/internal/test_ux.py
+++ b/tests/hikari/internal/test_ux.py
@@ -233,48 +233,36 @@ def test_when_not_allow_color(self):
         assert ux.supports_color(False, True) is False
 
     def test_when_CLICOLOR_FORCE_in_env(self):
-        with mock.patch.object(os, "getenv", return_value="1") as getenv:
+        with mock.patch.dict(os.environ, {"CLICOLOR_FORCE": "1"}, clear=True):
             assert ux.supports_color(True, False) is True
 
-        getenv.assert_called_once_with("CLICOLOR_FORCE", "0")
-
     def test_when_force_color(self):
-        with mock.patch.object(os, "getenv", return_value="0") as getenv:
+        with mock.patch.dict(os.environ, {"CLICOLOR_FORCE": "0"}, clear=True):
             assert ux.supports_color(True, True) is True
 
-        getenv.assert_called_once_with("CLICOLOR_FORCE", "0")
-
     def test_when_CLICOLOR_and_is_a_tty(self):
         with mock.patch.object(sys.stdout, "isatty", return_value=True):
-            with mock.patch.object(os, "getenv", side_effect=["0", "1"]) as getenv:
+            with mock.patch.dict(os.environ, {"CLICOLOR_FORCE": "0", "CLICOLOR": "1"}, clear=True):
                 assert ux.supports_color(True, False) is True
 
-        assert getenv.call_count == 2
-        getenv.assert_has_calls([mock.call("CLICOLOR_FORCE", "0"), mock.call("CLICOLOR", "0")])
+    def test_when_CLICOLOR_is_0(self):
+        with mock.patch.object(sys.stdout, "isatty", return_value=True):
+            with mock.patch.dict(os.environ, {"CLICOLOR_FORCE": "0", "CLICOLOR": "0"}, clear=True):
+                assert ux.supports_color(True, False) is False
 
     @pytest.mark.parametrize("colorterm", ["truecolor", "24bit", "TRUECOLOR", "24BIT"])
     def test_when_COLORTERM_has_correct_value(self, colorterm):
-        with mock.patch.object(os, "getenv", side_effect=["0", "0", colorterm]) as getenv:
+        with mock.patch.dict(os.environ, {"COLORTERM": colorterm}, clear=True):
             assert ux.supports_color(True, False) is True
 
-        assert getenv.call_count == 3
-        getenv.assert_has_calls(
-            [mock.call("CLICOLOR_FORCE", "0"), mock.call("CLICOLOR", "0"), mock.call("COLORTERM", "")]
-        )
-
     def test_when_plat_is_Pocket_PC(self):
         stack = contextlib.ExitStack()
-        getenv = stack.enter_context(mock.patch.object(os, "getenv", return_value="0"))
+        stack.enter_context(mock.patch.dict(os.environ, {}, clear=True))
         stack.enter_context(mock.patch.object(sys, "platform", new="Pocket PC"))
 
         with stack:
             assert ux.supports_color(True, False) is False
 
-        assert getenv.call_count == 3
-        getenv.assert_has_calls(
-            [mock.call("CLICOLOR_FORCE", "0"), mock.call("CLICOLOR", "0"), mock.call("COLORTERM", "")]
-        )
-
     @pytest.mark.parametrize(
         ("term_program", "ansicon", "isatty", "expected"),
         [
@@ -288,78 +276,49 @@ def test_when_plat_is_Pocket_PC(self):
         ],
     )
     def test_when_plat_is_win32(self, term_program, ansicon, isatty, expected):
+        environ = {"TERM_PROGRAM": term_program}
+        if ansicon:
+            environ["ANSICON"] = "ooga booga"
+
         stack = contextlib.ExitStack()
-        getenv = stack.enter_context(mock.patch.object(os, "getenv", side_effect=["0", "0", "", term_program, ""]))
+        stack.enter_context(mock.patch.dict(os.environ, environ, clear=True))
         stack.enter_context(mock.patch.object(sys.stdout, "isatty", return_value=isatty))
         stack.enter_context(mock.patch.object(sys, "platform", new="win32"))
-        stack.enter_context(mock.patch.object(os, "environ", new=["ANSICON"] if ansicon else []))
 
         with stack:
             assert ux.supports_color(True, False) is expected
 
-        assert getenv.call_count == 5
-        getenv.assert_has_calls(
-            [
-                mock.call("CLICOLOR_FORCE", "0"),
-                mock.call("CLICOLOR", "0"),
-                mock.call("COLORTERM", ""),
-                mock.call("TERM_PROGRAM", None),
-                mock.call("PYCHARM_HOSTED", ""),
-            ]
-        )
-
     @pytest.mark.parametrize("isatty", [True, False])
     def test_when_plat_is_not_win32(self, isatty):
         stack = contextlib.ExitStack()
-        getenv = stack.enter_context(mock.patch.object(os, "getenv", side_effect=["0", "0", "", ""]))
+        stack.enter_context(mock.patch.dict(os.environ, {}, clear=True))
         stack.enter_context(mock.patch.object(sys.stdout, "isatty", return_value=isatty))
         stack.enter_context(mock.patch.object(sys, "platform", new="linux"))
 
         with stack:
             assert ux.supports_color(True, False) is isatty
 
-        assert getenv.call_count == 4
-        getenv.assert_has_calls(
-            [
-                mock.call("CLICOLOR_FORCE", "0"),
-                mock.call("CLICOLOR", "0"),
-                mock.call("COLORTERM", ""),
-                mock.call("PYCHARM_HOSTED", ""),
-            ]
-        )
-
     @pytest.mark.parametrize("isatty", [True, False])
     @pytest.mark.parametrize("plat", ["linux", "win32"])
     def test_when_PYCHARM_HOSTED(self, isatty, plat):
         stack = contextlib.ExitStack()
-        getenv = stack.enter_context(mock.patch.object(os, "getenv", return_value="0"))
+        stack.enter_context(mock.patch.dict(os.environ, {"PYCHARM_HOSTED": "OOGA BOOGA"}, clear=True))
         stack.enter_context(mock.patch.object(sys.stdout, "isatty", return_value=isatty))
         stack.enter_context(mock.patch.object(sys, "platform", new=plat))
 
         with stack:
             assert ux.supports_color(True, False) is True
 
-        if plat == "win32":
-            assert getenv.call_count == 5
-            getenv.assert_has_calls(
-                [
-                    mock.call("CLICOLOR_FORCE", "0"),
-                    mock.call("CLICOLOR", "0"),
-                    mock.call("COLORTERM", ""),
-                    mock.call("TERM_PROGRAM", None),
-                    mock.call("PYCHARM_HOSTED", ""),
-                ]
-            )
-        else:
-            assert getenv.call_count == 4
-            getenv.assert_has_calls(
-                [
-                    mock.call("CLICOLOR_FORCE", "0"),
-                    mock.call("CLICOLOR", "0"),
-                    mock.call("COLORTERM", ""),
-                    mock.call("PYCHARM_HOSTED", ""),
-                ]
-            )
+    @pytest.mark.parametrize("isatty", [True, False])
+    @pytest.mark.parametrize("plat", ["linux", "win32"])
+    def test_when_WT_SESSION(self, isatty, plat):
+        stack = contextlib.ExitStack()
+        stack.enter_context(mock.patch.dict(os.environ, {"WT_SESSION": "OOGA BOOGA"}, clear=True))
+        stack.enter_context(mock.patch.object(sys.stdout, "isatty", return_value=isatty))
+        stack.enter_context(mock.patch.object(sys, "platform", new=plat))
+
+        with stack:
+            assert ux.supports_color(True, False) is True
 
 
 class TestHikariVersion: