Adapt tests to theming when executed with -j

Author: ambv

Lib/test/support/__init__.py

@@ -2855,36 +2855,59 @@ def is_slot_wrapper(name, value):
 
 
 @contextlib.contextmanager
-def no_color():
+def force_color(color: bool):
     import _colorize
     from .os_helper import EnvironmentVarGuard
 
     with (
-        swap_attr(_colorize, "can_colorize", lambda file=None: False),
+        swap_attr(_colorize, "can_colorize", lambda file=None: color),
         EnvironmentVarGuard() as env,
     ):
         env.unset("FORCE_COLOR", "NO_COLOR", "PYTHON_COLORS")
-        env.set("NO_COLOR", "1")
+        env.set("FORCE_COLOR" if color else "NO_COLOR", "1")
         yield
 
 
+def force_colorized(func):
+    """Force the terminal to be colorized."""
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        with force_color(True):
+            return func(*args, **kwargs)
+    return wrapper
+
+
 def force_not_colorized(func):
-    """Force the terminal not to be colorized."""
+    """Force the terminal NOT to be colorized."""
     @functools.wraps(func)
     def wrapper(*args, **kwargs):
-        with no_color():
+        with force_color(False):
             return func(*args, **kwargs)
     return wrapper
 
 
+def force_colorized_test_class(cls):
+    """Force the terminal to be colorized for the entire test class."""
+    original_setUpClass = cls.setUpClass
+
+    @classmethod
+    @functools.wraps(cls.setUpClass)
+    def new_setUpClass(cls):
+        cls.enterClassContext(force_color(True))
+        original_setUpClass()
+
+    cls.setUpClass = new_setUpClass
+    return cls
+
+
 def force_not_colorized_test_class(cls):
-    """Force the terminal not to be colorized for the entire test class."""
+    """Force the terminal NOT to be colorized for the entire test class."""
     original_setUpClass = cls.setUpClass
 
     @classmethod
     @functools.wraps(cls.setUpClass)
     def new_setUpClass(cls):
-        cls.enterClassContext(no_color())
+        cls.enterClassContext(force_color(False))
         original_setUpClass()
 
     cls.setUpClass = new_setUpClass

Lib/test/test_pyrepl/support.py

@@ -113,9 +113,6 @@ def handle_all_events(
     prepare_console=partial(prepare_console, width=10),
 )
 
-reader_no_colors = partial(prepare_reader, can_colorize=False)
-reader_force_colors = partial(prepare_reader, can_colorize=True)
-
 
 class FakeConsole(Console):
     def __init__(self, events, encoding="utf-8") -> None:

Lib/test/test_pyrepl/test_reader.py

@@ -4,11 +4,11 @@
 from textwrap import dedent
 from unittest import TestCase
 from unittest.mock import MagicMock
+from test.support import force_colorized_test_class, force_not_colorized_test_class
 
 from .support import handle_all_events, handle_events_narrow_console
 from .support import ScreenEqualMixin, code_to_events
-from .support import prepare_console, reader_force_colors
-from .support import reader_no_colors as prepare_reader
+from .support import prepare_reader, prepare_console
 from _pyrepl.console import Event
 from _pyrepl.reader import Reader
 from _colorize import default_theme
@@ -18,6 +18,7 @@
 colors = {overrides.get(k, k[0].lower()): v for k, v in default_theme.repl.items()}
 
 
+@force_not_colorized_test_class
 class TestReader(ScreenEqualMixin, TestCase):
     def test_calc_screen_wrap_simple(self):
         events = code_to_events(10 * "a")
@@ -127,13 +128,6 @@ def test_setpos_for_xy_simple(self):
         reader.setpos_from_xy(0, 0)
         self.assertEqual(reader.pos, 0)
 
-    def test_control_characters(self):
-        code = 'flag = "🏳️‍🌈"'
-        events = code_to_events(code)
-        reader, _ = handle_all_events(events, prepare_reader=reader_force_colors)
-        self.assert_screen_equal(reader, 'flag = "🏳️\\u200d🌈"', clean=True)
-        self.assert_screen_equal(reader, 'flag {o}={z} {s}"🏳️\\u200d🌈"{z}'.format(**colors))
-
     def test_setpos_from_xy_multiple_lines(self):
         # fmt: off
         code = (
@@ -364,6 +358,8 @@ def test_setpos_from_xy_for_non_printing_char(self):
         reader.setpos_from_xy(8, 0)
         self.assertEqual(reader.pos, 7)
 
+@force_colorized_test_class
+class TestReaderInColor(ScreenEqualMixin, TestCase):
     def test_syntax_highlighting_basic(self):
         code = dedent(
             """\
@@ -403,7 +399,7 @@ def funct(case: str = sys.platform) -> None:
         )
         expected_sync = expected.format(a="", **colors)
         events = code_to_events(code)
-        reader, _ = handle_all_events(events, prepare_reader=reader_force_colors)
+        reader, _ = handle_all_events(events)
         self.assert_screen_equal(reader, code, clean=True)
         self.assert_screen_equal(reader, expected_sync)
         self.assertEqual(reader.pos, 2**7 + 2**8)
@@ -416,7 +412,7 @@ def funct(case: str = sys.platform) -> None:
             [Event(evt="key", data="up", raw=bytearray(b"\x1bOA"))] * 13,
             code_to_events("async "),
         )
-        reader, _ = handle_all_events(more_events, prepare_reader=reader_force_colors)
+        reader, _ = handle_all_events(more_events)
         self.assert_screen_equal(reader, expected_async)
         self.assertEqual(reader.pos, 21)
         self.assertEqual(reader.cxy, (6, 1))
@@ -433,7 +429,7 @@ def unfinished_function(arg: str = "still typing
             """
         ).format(**colors)
         events = code_to_events(code)
-        reader, _ = handle_all_events(events, prepare_reader=reader_force_colors)
+        reader, _ = handle_all_events(events)
         self.assert_screen_equal(reader, code, clean=True)
         self.assert_screen_equal(reader, expected)
 
@@ -451,7 +447,7 @@ def unfinished_function(
             """
         ).format(**colors)
         events = code_to_events(code)
-        reader, _ = handle_all_events(events, prepare_reader=reader_force_colors)
+        reader, _ = handle_all_events(events)
         self.assert_screen_equal(reader, code, clean=True)
         self.assert_screen_equal(reader, expected)
 
@@ -471,7 +467,7 @@ def unfinished_function():
             """
         ).format(**colors)
         events = code_to_events(code)
-        reader, _ = handle_all_events(events, prepare_reader=reader_force_colors)
+        reader, _ = handle_all_events(events)
         self.assert_screen_equal(reader, code, clean=True)
         self.assert_screen_equal(reader, expected)
 
@@ -497,6 +493,13 @@ def unfinished_function():
             """
         ).format(OB="{", CB="}", **colors)
         events = code_to_events(code)
-        reader, _ = handle_all_events(events, prepare_reader=reader_force_colors)
+        reader, _ = handle_all_events(events)
         self.assert_screen_equal(reader, code, clean=True)
         self.assert_screen_equal(reader, expected)
+
+    def test_control_characters(self):
+        code = 'flag = "🏳️‍🌈"'
+        events = code_to_events(code)
+        reader, _ = handle_all_events(events)
+        self.assert_screen_equal(reader, 'flag = "🏳️\\u200d🌈"', clean=True)
+        self.assert_screen_equal(reader, 'flag {o}={z} {s}"🏳️\\u200d🌈"{z}'.format(**colors))

Lib/test/test_pyrepl/test_unix_console.py

@@ -3,11 +3,12 @@
 import sys
 import unittest
 from functools import partial
-from test.support import os_helper
+from test.support import os_helper, force_not_colorized_test_class
+
 from unittest import TestCase
 from unittest.mock import MagicMock, call, patch, ANY
 
-from .support import handle_all_events, code_to_events, reader_no_colors
+from .support import handle_all_events, code_to_events
 
 try:
     from _pyrepl.console import Event
@@ -33,12 +34,10 @@ def unix_console(events, **kwargs):
 
 handle_events_unix_console = partial(
     handle_all_events,
-    prepare_reader=reader_no_colors,
     prepare_console=unix_console,
 )
 handle_events_narrow_unix_console = partial(
     handle_all_events,
-    prepare_reader=reader_no_colors,
     prepare_console=partial(unix_console, width=5),
 )
 handle_events_short_unix_console = partial(
@@ -120,6 +119,7 @@ def unix_console(events, **kwargs):
 )
 @patch("termios.tcsetattr", lambda a, b, c: None)
 @patch("os.write")
+@force_not_colorized_test_class
 class TestConsole(TestCase):
     def test_simple_addition(self, _os_write):
         code = "12+34"
@@ -255,9 +255,7 @@ def test_resize_bigger_on_multiline_function(self, _os_write):
         # fmt: on
 
         events = itertools.chain(code_to_events(code))
-        reader, console = handle_events_short_unix_console(
-            events, prepare_reader=reader_no_colors
-        )
+        reader, console = handle_events_short_unix_console(events)
 
         console.height = 2
         console.getheightwidth = MagicMock(lambda _: (2, 80))

Lib/test/test_pyrepl/test_windows_console.py

@@ -7,12 +7,13 @@
 
 import itertools
 from functools import partial
+from test.support import force_not_colorized_test_class
 from typing import Iterable
 from unittest import TestCase
 from unittest.mock import MagicMock, call
 
 from .support import handle_all_events, code_to_events
-from .support import reader_no_colors as default_prepare_reader
+from .support import prepare_reader as default_prepare_reader
 
 try:
     from _pyrepl.console import Event, Console
@@ -28,6 +29,7 @@
     pass
 
 
+@force_not_colorized_test_class
 class WindowsConsoleTests(TestCase):
     def console(self, events, **kwargs) -> Console:
         console = WindowsConsole()