From 44a6efdafa702c033f812381ff7d3e9990979866 Mon Sep 17 00:00:00 2001
From: Victor Stinner <vstinner@python.org>
Date: Wed, 6 Sep 2023 17:34:31 +0200
Subject: [PATCH] gh-108851: Fix tomllib recursion tests (#108853)

* Add get_recursion_available() and get_recursion_depth() functions
  to the test.support module.
* Change infinite_recursion() default max_depth from 75 to 100.
* Fix test_tomllib recursion tests for WASI buildbots: reduce the
  recursion limit and compute the maximum nested array/dict depending
  on the current available recursion limit.
* test.pythoninfo logs sys.getrecursionlimit().
* Enhance test_sys tests on sys.getrecursionlimit()
  and sys.setrecursionlimit().

Backport notes:

* Set support.infinite_recursion() minimum to 4 frames.
* test_support.test_get_recursion_depth() uses limit-2, apparently
  f-string counts for 2 frames in Python 3.11.
* test_sys.test_setrecursionlimit_to_depth() tests depth+2 instead of
  depth+1.

(cherry picked from commit 8ff11425783806f8cb78e99f667546b1f7f3428e)
---
 Lib/test/pythoninfo.py                        |  1 +
 Lib/test/support/__init__.py                  | 57 +++++++++++--
 Lib/test/test_support.py                      | 79 +++++++++++++++++++
 Lib/test/test_sys.py                          | 57 +++++++------
 Lib/test/test_tomllib/test_misc.py            | 27 +++++--
 ...-09-03-21-18-35.gh-issue-108851.CCuHyI.rst |  2 +
 ...-09-03-21-41-10.gh-issue-108851.xFTYOE.rst |  3 +
 7 files changed, 184 insertions(+), 42 deletions(-)
 create mode 100644 Misc/NEWS.d/next/Tests/2023-09-03-21-18-35.gh-issue-108851.CCuHyI.rst
 create mode 100644 Misc/NEWS.d/next/Tests/2023-09-03-21-41-10.gh-issue-108851.xFTYOE.rst

diff --git a/Lib/test/pythoninfo.py b/Lib/test/pythoninfo.py
index 3240e2e3bca80d..2575a9ae9e5538 100644
--- a/Lib/test/pythoninfo.py
+++ b/Lib/test/pythoninfo.py
@@ -112,6 +112,7 @@ def collect_sys(info_add):
 
     call_func(info_add, 'sys.androidapilevel', sys, 'getandroidapilevel')
     call_func(info_add, 'sys.windowsversion', sys, 'getwindowsversion')
+    call_func(info_add, 'sys.getrecursionlimit', sys, 'getrecursionlimit')
 
     encoding = sys.getfilesystemencoding()
     if hasattr(sys, 'getfilesystemencodeerrors'):
diff --git a/Lib/test/support/__init__.py b/Lib/test/support/__init__.py
index 94967f2c302a8c..05b33753724dbc 100644
--- a/Lib/test/support/__init__.py
+++ b/Lib/test/support/__init__.py
@@ -2197,20 +2197,61 @@ def check_disallow_instantiation(testcase, tp, *args, **kwds):
     msg = f"cannot create '{re.escape(qualname)}' instances"
     testcase.assertRaisesRegex(TypeError, msg, tp, *args, **kwds)
 
+def get_recursion_depth():
+    """Get the recursion depth of the caller function.
+
+    In the __main__ module, at the module level, it should be 1.
+    """
+    try:
+        import _testinternalcapi
+        depth = _testinternalcapi.get_recursion_depth()
+    except (ImportError, RecursionError) as exc:
+        # sys._getframe() + frame.f_back implementation.
+        try:
+            depth = 0
+            frame = sys._getframe()
+            while frame is not None:
+                depth += 1
+                frame = frame.f_back
+        finally:
+            # Break any reference cycles.
+            frame = None
+
+    # Ignore get_recursion_depth() frame.
+    return max(depth - 1, 1)
+
+def get_recursion_available():
+    """Get the number of available frames before RecursionError.
+
+    It depends on the current recursion depth of the caller function and
+    sys.getrecursionlimit().
+    """
+    limit = sys.getrecursionlimit()
+    depth = get_recursion_depth()
+    return limit - depth
+
 @contextlib.contextmanager
-def infinite_recursion(max_depth=75):
+def set_recursion_limit(limit):
+    """Temporarily change the recursion limit."""
+    original_limit = sys.getrecursionlimit()
+    try:
+        sys.setrecursionlimit(limit)
+        yield
+    finally:
+        sys.setrecursionlimit(original_limit)
+
+def infinite_recursion(max_depth=100):
     """Set a lower limit for tests that interact with infinite recursions
     (e.g test_ast.ASTHelpers_Test.test_recursion_direct) since on some
     debug windows builds, due to not enough functions being inlined the
     stack size might not handle the default recursion limit (1000). See
     bpo-11105 for details."""
-
-    original_depth = sys.getrecursionlimit()
-    try:
-        sys.setrecursionlimit(max_depth)
-        yield
-    finally:
-        sys.setrecursionlimit(original_depth)
+    if max_depth < 4:
+        raise ValueError("max_depth must be at least 4, got {max_depth}")
+    depth = get_recursion_depth()
+    depth = max(depth - 1, 1)  # Ignore infinite_recursion() frame.
+    limit = depth + max_depth
+    return set_recursion_limit(limit)
 
 def ignore_deprecations_from(module: str, *, like: str) -> object:
     token = object()
diff --git a/Lib/test/test_support.py b/Lib/test/test_support.py
index 2a33889b7e97ac..54e532064740f3 100644
--- a/Lib/test/test_support.py
+++ b/Lib/test/test_support.py
@@ -698,6 +698,85 @@ def test_has_strftime_extensions(self):
         else:
             self.assertTrue(support.has_strftime_extensions)
 
+    def test_get_recursion_depth(self):
+        # test support.get_recursion_depth()
+        code = textwrap.dedent("""
+            from test import support
+            import sys
+
+            def check(cond):
+                if not cond:
+                    raise AssertionError("test failed")
+
+            # depth 1
+            check(support.get_recursion_depth() == 1)
+
+            # depth 2
+            def test_func():
+                check(support.get_recursion_depth() == 2)
+            test_func()
+
+            def test_recursive(depth, limit):
+                if depth >= limit:
+                    # cannot call get_recursion_depth() at this depth,
+                    # it can raise RecursionError
+                    return
+                get_depth = support.get_recursion_depth()
+                print(f"test_recursive: {depth}/{limit}: "
+                      f"get_recursion_depth() says {get_depth}")
+                check(get_depth == depth)
+                test_recursive(depth + 1, limit)
+
+            # depth up to 25
+            with support.infinite_recursion(max_depth=25):
+                limit = sys.getrecursionlimit()
+                print(f"test with sys.getrecursionlimit()={limit}")
+                # Use limit-2 since f-string seems to consume 2 frames.
+                test_recursive(2, limit - 2)
+
+            # depth up to 500
+            with support.infinite_recursion(max_depth=500):
+                limit = sys.getrecursionlimit()
+                print(f"test with sys.getrecursionlimit()={limit}")
+                # limit-2 since f-string seems to consume 2 frames
+                test_recursive(2, limit - 2)
+        """)
+        script_helper.assert_python_ok("-c", code)
+
+    def test_recursion(self):
+        # Test infinite_recursion() and get_recursion_available() functions.
+        def recursive_function(depth):
+            if depth:
+                recursive_function(depth - 1)
+
+        for max_depth in (5, 25, 250):
+            with support.infinite_recursion(max_depth):
+                available = support.get_recursion_available()
+
+                # Recursion up to 'available' additional frames should be OK.
+                recursive_function(available)
+
+                # Recursion up to 'available+1' additional frames must raise
+                # RecursionError. Avoid self.assertRaises(RecursionError) which
+                # can consume more than 3 frames and so raises RecursionError.
+                try:
+                    recursive_function(available + 1)
+                except RecursionError:
+                    pass
+                else:
+                    self.fail("RecursionError was not raised")
+
+        # Test the bare minimumum: max_depth=4
+        with support.infinite_recursion(4):
+            try:
+                recursive_function(4)
+            except RecursionError:
+                pass
+            else:
+                self.fail("RecursionError was not raised")
+
+        #self.assertEqual(available, 2)
+
     # XXX -follows a list of untested API
     # make_legacy_pyc
     # is_resource_enabled
diff --git a/Lib/test/test_sys.py b/Lib/test/test_sys.py
index 6f56c9ef97e7c3..86cf1a794f973c 100644
--- a/Lib/test/test_sys.py
+++ b/Lib/test/test_sys.py
@@ -269,20 +269,29 @@ def test_switchinterval(self):
         finally:
             sys.setswitchinterval(orig)
 
-    def test_recursionlimit(self):
+    def test_getrecursionlimit(self):
+        limit = sys.getrecursionlimit()
+        self.assertIsInstance(limit, int)
+        self.assertGreater(limit, 1)
+
         self.assertRaises(TypeError, sys.getrecursionlimit, 42)
-        oldlimit = sys.getrecursionlimit()
-        self.assertRaises(TypeError, sys.setrecursionlimit)
-        self.assertRaises(ValueError, sys.setrecursionlimit, -42)
-        sys.setrecursionlimit(10000)
-        self.assertEqual(sys.getrecursionlimit(), 10000)
-        sys.setrecursionlimit(oldlimit)
+
+    def test_setrecursionlimit(self):
+        old_limit = sys.getrecursionlimit()
+        try:
+            sys.setrecursionlimit(10_005)
+            self.assertEqual(sys.getrecursionlimit(), 10_005)
+
+            self.assertRaises(TypeError, sys.setrecursionlimit)
+            self.assertRaises(ValueError, sys.setrecursionlimit, -42)
+        finally:
+            sys.setrecursionlimit(old_limit)
 
     def test_recursionlimit_recovery(self):
         if hasattr(sys, 'gettrace') and sys.gettrace():
             self.skipTest('fatal error if run with a trace function')
 
-        oldlimit = sys.getrecursionlimit()
+        old_limit = sys.getrecursionlimit()
         def f():
             f()
         try:
@@ -301,35 +310,31 @@ def f():
                 with self.assertRaises(RecursionError):
                     f()
         finally:
-            sys.setrecursionlimit(oldlimit)
+            sys.setrecursionlimit(old_limit)
 
     @test.support.cpython_only
-    def test_setrecursionlimit_recursion_depth(self):
+    def test_setrecursionlimit_to_depth(self):
         # Issue #25274: Setting a low recursion limit must be blocked if the
         # current recursion depth is already higher than limit.
 
-        from _testinternalcapi import get_recursion_depth
-
-        def set_recursion_limit_at_depth(depth, limit):
-            recursion_depth = get_recursion_depth()
-            if recursion_depth >= depth:
+        old_limit = sys.getrecursionlimit()
+        try:
+            depth = support.get_recursion_depth()
+            with self.subTest(limit=sys.getrecursionlimit(), depth=depth):
+                # depth + 2 is OK
+                sys.setrecursionlimit(depth + 2)
+
+                # reset the limit to be able to call self.assertRaises()
+                # context manager
+                sys.setrecursionlimit(old_limit)
                 with self.assertRaises(RecursionError) as cm:
-                    sys.setrecursionlimit(limit)
+                    sys.setrecursionlimit(depth + 1)
                 self.assertRegex(str(cm.exception),
                                  "cannot set the recursion limit to [0-9]+ "
                                  "at the recursion depth [0-9]+: "
                                  "the limit is too low")
-            else:
-                set_recursion_limit_at_depth(depth, limit)
-
-        oldlimit = sys.getrecursionlimit()
-        try:
-            sys.setrecursionlimit(1000)
-
-            for limit in (10, 25, 50, 75, 100, 150, 200):
-                set_recursion_limit_at_depth(limit, limit)
         finally:
-            sys.setrecursionlimit(oldlimit)
+            sys.setrecursionlimit(old_limit)
 
     def test_getwindowsversion(self):
         # Raise SkipTest if sys doesn't have getwindowsversion attribute
diff --git a/Lib/test/test_tomllib/test_misc.py b/Lib/test/test_tomllib/test_misc.py
index a477a219fd9ebd..9e677a337a2835 100644
--- a/Lib/test/test_tomllib/test_misc.py
+++ b/Lib/test/test_tomllib/test_misc.py
@@ -9,6 +9,7 @@
 import sys
 import tempfile
 import unittest
+from test import support
 
 from . import tomllib
 
@@ -92,13 +93,23 @@ def test_deepcopy(self):
         self.assertEqual(obj_copy, expected_obj)
 
     def test_inline_array_recursion_limit(self):
-        # 465 with default recursion limit
-        nest_count = int(sys.getrecursionlimit() * 0.465)
-        recursive_array_toml = "arr = " + nest_count * "[" + nest_count * "]"
-        tomllib.loads(recursive_array_toml)
+        with support.infinite_recursion(max_depth=100):
+            available = support.get_recursion_available()
+            nest_count = (available // 2) - 2
+            # Add details if the test fails
+            with self.subTest(limit=sys.getrecursionlimit(),
+                              available=available,
+                              nest_count=nest_count):
+                recursive_array_toml = "arr = " + nest_count * "[" + nest_count * "]"
+                tomllib.loads(recursive_array_toml)
 
     def test_inline_table_recursion_limit(self):
-        # 310 with default recursion limit
-        nest_count = int(sys.getrecursionlimit() * 0.31)
-        recursive_table_toml = nest_count * "key = {" + nest_count * "}"
-        tomllib.loads(recursive_table_toml)
+        with support.infinite_recursion(max_depth=100):
+            available = support.get_recursion_available()
+            nest_count = (available // 3) - 1
+            # Add details if the test fails
+            with self.subTest(limit=sys.getrecursionlimit(),
+                              available=available,
+                              nest_count=nest_count):
+                recursive_table_toml = nest_count * "key = {" + nest_count * "}"
+                tomllib.loads(recursive_table_toml)
diff --git a/Misc/NEWS.d/next/Tests/2023-09-03-21-18-35.gh-issue-108851.CCuHyI.rst b/Misc/NEWS.d/next/Tests/2023-09-03-21-18-35.gh-issue-108851.CCuHyI.rst
new file mode 100644
index 00000000000000..7a5b3052af22f2
--- /dev/null
+++ b/Misc/NEWS.d/next/Tests/2023-09-03-21-18-35.gh-issue-108851.CCuHyI.rst
@@ -0,0 +1,2 @@
+Add ``get_recursion_available()`` and ``get_recursion_depth()`` functions to
+the :mod:`test.support` module. Patch by Victor Stinner.
diff --git a/Misc/NEWS.d/next/Tests/2023-09-03-21-41-10.gh-issue-108851.xFTYOE.rst b/Misc/NEWS.d/next/Tests/2023-09-03-21-41-10.gh-issue-108851.xFTYOE.rst
new file mode 100644
index 00000000000000..b35aaebb410afb
--- /dev/null
+++ b/Misc/NEWS.d/next/Tests/2023-09-03-21-41-10.gh-issue-108851.xFTYOE.rst
@@ -0,0 +1,3 @@
+Fix ``test_tomllib`` recursion tests for WASI buildbots: reduce the recursion
+limit and compute the maximum nested array/dict depending on the current
+available recursion limit. Patch by Victor Stinner.