Respect the recursion limit in nondeterministic mode (#1)

* Respect the recursion limit in nondeterministic mode

* Fix depth for children of children

Author: jg-rp

jsonpath_rfc9535/segments.py

@@ -76,7 +76,7 @@ def resolve(self, nodes: Iterable[JSONPathNode]) -> Iterable[JSONPathNode]:
         # XXX: This feels like a bit of a hack.
         visitor = (
             self._nondeterministic_visit
-            if self.env.nondeterministic and random.random() < 0.2  # noqa: S311, PLR2004
+            if self.env.nondeterministic and random.random() < 0.8  # noqa: S311, PLR2004
             else self._visit
         )
 
@@ -86,7 +86,7 @@ def resolve(self, nodes: Iterable[JSONPathNode]) -> Iterable[JSONPathNode]:
                     yield from selector.resolve(_node)
 
     def _visit(self, node: JSONPathNode, depth: int = 1) -> Iterable[JSONPathNode]:
-        """Pre order node traversal."""
+        """Depth-first, pre-order node traversal."""
         if depth > self.env.max_recursion_depth:
             raise JSONPathRecursionError("recursion limit exceeded", token=self.token)
 
@@ -136,20 +136,31 @@ def _children(node: JSONPathNode) -> Iterable[JSONPathNode]:
                             root=node.root,
                         )
 
-        queue: Deque[JSONPathNode] = deque(_children(root))
-        yield root
+        # (node, depth) tuples
+        queue: Deque[Tuple[JSONPathNode, int]] = deque()
+
+        yield root  # visit the root node
+        queue.extend([(child, 1) for child in _children(root)])  # queue root's children
 
         while queue:
-            _node = queue.popleft()
+            _node, depth = queue.popleft()
+
+            if depth >= self.env.max_recursion_depth:
+                raise JSONPathRecursionError(
+                    "recursion limit exceeded", token=self.token
+                )
+
             yield _node
+
             # Visit child nodes now or queue them for later?
             visit_children = random.choice([True, False])  # noqa: S311
+
             for child in _children(_node):
                 if visit_children:
                     yield child
-                    queue.extend(_children(child))
+                    queue.extend([(child, depth + 2) for child in _children(child)])
                 else:
-                    queue.append(child)
+                    queue.append((child, depth + 1))
 
     def __str__(self) -> str:
         return f"..[{', '.join(str(itm) for itm in self.selectors)}]"

tests/test_errors.py

@@ -64,11 +64,26 @@ class MockEnv(JSONPathEnvironment):
         env.find(query, data)
 
 
+def test_recursive_data_nondeterministic() -> None:
+    class MockEnv(JSONPathEnvironment):
+        nondeterministic = True
+
+    env = MockEnv()
+    query = "$..a"
+    arr: List[Any] = []
+    data: Any = {"foo": arr}
+    arr.append(data)
+
+    with pytest.raises(JSONPathRecursionError):
+        env.find(query, data)
+
+
 class FilterLiteralTestCase(NamedTuple):
     description: str
     query: str
 
 
+# TODO: add these to the CTS?
 BAD_FILTER_LITERAL_TEST_CASES: List[FilterLiteralTestCase] = [
     FilterLiteralTestCase("just true", "$[?true]"),
     FilterLiteralTestCase("just string", "$[?'foo']"),

tests/test_nondeterminism.py

@@ -10,6 +10,7 @@
 from typing import Any
 from typing import List
 from typing import Optional
+from typing import Tuple
 
 import pytest
 
@@ -37,12 +38,16 @@ def valid_cases() -> List[Case]:
     return [case for case in cases() if not case.invalid_selector]
 
 
+def nondeterministic_cases() -> List[Case]:
+    return [case for case in valid_cases() if isinstance(case.results, list)]
+
+
 class MockEnv(JSONPathEnvironment):
     nondeterministic = True
 
 
 @pytest.mark.parametrize("case", valid_cases(), ids=operator.attrgetter("name"))
-def test_nondeterminism(case: Case) -> None:
+def test_nondeterminism_valid_cases(case: Case) -> None:
     assert case.document is not None
     env = MockEnv()
     rv = env.find(case.selector, case.document).values()
@@ -51,3 +56,28 @@ def test_nondeterminism(case: Case) -> None:
         assert rv in case.results
     else:
         assert rv == case.result
+
+
+@pytest.mark.parametrize(
+    "case", nondeterministic_cases(), ids=operator.attrgetter("name")
+)
+def test_nondeterminism(case: Case) -> None:
+    """Test that we agree with CTS when it comes to nondeterministic results."""
+    assert case.document is not None
+    assert case.results is not None
+
+    def _result_repr(rv: List[object]) -> Tuple[str, ...]:
+        """Return a hashable representation of a result list."""
+        return tuple([str(value) for value in rv])
+
+    env = MockEnv()
+
+    # Repeat enough times to has high probability that we've covered all
+    # valid permutations.
+    results = {
+        _result_repr(env.find(case.selector, case.document).values())
+        for _ in range(1000)
+    }
+
+    assert len(results) == len(case.results)
+    assert results == {_result_repr(result) for result in case.results}