Use setUpClass/tearDownClass instead of setUp/tearDown

python · hugovk · Jan 13, 2025 · Dec 11, 2024 · Dec 11, 2024 · Jan 9, 2025
commit e9aa442e13aa5fc13219ce963380c3866911ae06
diff --git a/Lib/test/support/__init__.py b/Lib/test/support/__init__.py
@@ -2870,24 +2870,30 @@ def wrapper(*args, **kwargs):
             _re_enable_terminal_color(original_fn, variables)
     return wrapper
 
-def force_not_colorized_test_class(cls):
-    """Force the terminal not to be colorized."""
-    original_setup = cls.setUp
-    original_teardown = cls.tearDown
-
-    @functools.wraps(cls.setUp)
-    def setUp_wrapper(self, *args, **kwargs):
-        self._original_fn, self._variables = _disable_terminal_color()
 
-        return original_setup(self, *args, **kwargs)
-
-    @functools.wraps(cls.tearDown)
-    def tearDown_wrapper(self, *args, **kwargs):
-        _re_enable_terminal_color(self._original_fn, self._variables)
-        return original_teardown(self, *args, **kwargs)
-
-    cls.setUp = setUp_wrapper
-    cls.tearDown = tearDown_wrapper
+def force_not_colorized_test_class(cls):
+    """Force the terminal not to be colorized for the entire test class."""
+    original_setUpClass = cls.setUpClass
+    original_tearDownClass = cls.tearDownClass
+
+    @classmethod
+    @functools.wraps(cls.setUpClass)
+    def new_setUpClass(cls):
+        original_fn, variables = _disable_terminal_color()
+        cls._original_fn = original_fn
+        cls._variables = variables
+        if original_setUpClass:
+            original_setUpClass()
+
+    @classmethod
+    @functools.wraps(cls.tearDownClass)
+    def new_tearDownClass(cls):
+        if original_tearDownClass:
+            original_tearDownClass()
+        _re_enable_terminal_color(cls._original_fn, cls._variables)
+
+    cls.setUpClass = new_setUpClass
+    cls.tearDownClass = new_tearDownClass
     return cls