python · sobolevn · May 28, 2024
diff --git a/Lib/test/libregrtest/refleak.py b/Lib/test/libregrtest/refleak.py
@@ -173,6 +173,8 @@ def get_pooled_int(value):
         interned_before = interned_after
 
         restore_support_xml(xml_filename)
+        refleak_helper._refleak_iteration += 1
+    refleak_helper._refleak_iteration = 0
 
     if not quiet:
         print(file=sys.stderr)

diff --git a/Lib/test/support/__init__.py b/Lib/test/support/__init__.py
@@ -1196,16 +1196,28 @@ def no_rerun(reason):
     impossible to find reference leaks. Provide a clear reason for skipping the
     test using the 'reason' parameter.
     """
-    def deco(func):
-        _has_run = False
-        def wrapper(self):
-            nonlocal _has_run
-            if _has_run:
-                self.skipTest(reason)
-            func(self)
-            _has_run = True
+    def decorator(func_or_class):
+        if isinstance(func_or_class, type):
+            setUpClass = func_or_class.__dict__.get('setUpClass')
+            if setUpClass is None:
+                def setUpClass(cls):
+                    super(func_or_class, cls).setUpClass()
+                setUpClass.__qualname__ = func_or_class.__qualname__ + '.setUpClass'
+                setUpClass.__module__ = func_or_class.__module__
+            else:
+                setUpClass = setUpClass.__func__
+            setUpClass = classmethod(decorator(setUpClass))
+            func_or_class.setUpClass = setUpClass
+            return func_or_class
+
+        @functools.wraps(func_or_class)
+        def wrapper(*args, **kwargs):
+            from test.support import refleak_helper
+            if refleak_helper._refleak_iteration > 0:
+                raise unittest.SkipTest(reason)
+            return func_or_class(*args, **kwargs)
         return wrapper
-    return deco
+    return decorator
 
 
 def refcount_test(test):

diff --git a/Lib/test/support/refleak_helper.py b/Lib/test/support/refleak_helper.py
@@ -6,3 +6,5 @@
 _hunting_for_refleaks = False
 def hunting_for_refleaks():
     return _hunting_for_refleaks
+
+_refleak_iteration = 0
@@ -39,16 +39,22 @@ def load_tests(loader, tests, pattern):
         for cls in test_classes:
             cls.__name__ += suffix
             cls.__qualname__ += suffix
+            orig_setUpClass = getattr(cls, 'setUpClass', None)
+            orig_tearDownClass = getattr(cls, 'tearDownClass', None)
             @classmethod
-            def setUpClass(cls_, module=module):
+            def setUpClass(cls_, module=module, orig_setUpClass=orig_setUpClass):
                 cls_._save_sys_modules = sys.modules.copy()
                 sys.modules[TESTS] = module
                 sys.modules['datetime'] = module.datetime_module
                 if hasattr(module, '_pydatetime'):
                     sys.modules['_pydatetime'] = module._pydatetime
                 sys.modules['_strptime'] = module._strptime
+                if orig_setUpClass is not None:
+                    orig_setUpClass()
             @classmethod
-            def tearDownClass(cls_):
+            def tearDownClass(cls_, orig_tearDownClass=orig_tearDownClass):
+                if orig_tearDownClass is not None:
+                    orig_tearDownClass()
                 sys.modules.clear()
                 sys.modules.update(cls_._save_sys_modules)
             cls.setUpClass = setUpClass