From 54c8a106c621f6f872ac820e7a8ecc0b4fce8531 Mon Sep 17 00:00:00 2001
From: Marko Ristin <marko.ristin@gmail.com>
Date: Tue, 6 Jul 2021 12:52:06 +0200
Subject: [PATCH] Add support for recursive data structures

We need to update the local namespace when we resolve type annotations
of recursive data structures when registering them over `DBC` meta-class
since they are still not available in the scope at that time.

This patch simply puts the class to be registered in the local
namespace.
---
 icontract_hypothesis/__init__.py              | 10 ++++--
 .../test_strategy_inference.py                | 36 +++++++++++++++++++
 2 files changed, 44 insertions(+), 2 deletions(-)

diff --git a/icontract_hypothesis/__init__.py b/icontract_hypothesis/__init__.py
index 96ace1d..760db2f 100644
--- a/icontract_hypothesis/__init__.py
+++ b/icontract_hypothesis/__init__.py
@@ -691,7 +691,12 @@ def _strategy_for_type(
     init = getattr(a_type, "__init__")
 
     if inspect.isfunction(init):
-        strategy = infer_strategy(init)
+        # Add a local namespace in case there are forward references.
+        #
+        # This is needed if we register a class through the ``icontract.DBCMeta``
+        # meta-class where it references itself. For example, a node in a linked list.
+        strategy = infer_strategy(init, localns={a_type.__name__: a_type})
+
     elif isinstance(init, icontract._checkers._SLOT_WRAPPER_TYPE):
         # We have to distinguish this special case which is used by named tuples and
         # possibly other optimized data structures.
@@ -1477,7 +1482,8 @@ def _register_with_hypothesis(cls: Type[T]) -> None:
         return
 
     if cls not in hypothesis.strategies._internal.types._global_type_lookup:
-        hypothesis.strategies.register_type_strategy(cls, _strategy_for_type(cls))
+        strategy = _strategy_for_type(cls)
+        hypothesis.strategies.register_type_strategy(custom_type=cls, strategy=strategy)
 
 
 def _hook_into_icontract_and_hypothesis() -> None:
diff --git a/tests/strategy_inference/test_strategy_inference.py b/tests/strategy_inference/test_strategy_inference.py
index f2f24bd..96fdc0c 100644
--- a/tests/strategy_inference/test_strategy_inference.py
+++ b/tests/strategy_inference/test_strategy_inference.py
@@ -253,6 +253,24 @@ def some_func(x: int, y: int) -> None:
         icontract_hypothesis.test_with_inferred_strategy(some_func)
 
 
+class SomeCyclicalGlobalClass(icontract.DBC):
+    """
+    Represent a class which has a cyclical dependency on itself.
+
+    For example, a node of a linked list.
+    """
+
+    value: int
+    next_node: Optional["SomeCyclicalGlobalClass"]
+
+    @icontract.require(lambda value: value > 0)
+    def __init__(
+        self, value: int, next_node: Optional["SomeCyclicalGlobalClass"]
+    ) -> None:
+        self.value = value
+        self.next_node = next_node
+
+
 # noinspection PyUnusedLocal
 class TestWithInferredStrategiesOnClasses(unittest.TestCase):
     def test_no_preconditions_and_no_argument_init(self) -> None:
@@ -661,6 +679,24 @@ def some_func(a_or_b: Union[A, B]) -> None:
 
         icontract_hypothesis.test_with_inferred_strategy(some_func)
 
+    def test_cyclical_data_structure(self) -> None:
+        def some_func(cyclical: SomeCyclicalGlobalClass) -> None:
+            pass
+
+        strategy = icontract_hypothesis.infer_strategy(some_func)
+
+        self.assertEqual(
+            "fixed_dictionaries({"
+            "'cyclical': fixed_dictionaries({"
+            "'next_node': one_of(none(), builds(SomeCyclicalGlobalClass)),\n"
+            "  'value': integers(min_value=1)}).map(lambda d: SomeCyclicalGlobalClass(**d))})",
+            str(strategy),
+        )
+
+        # We can not execute this strategy, as ``builds`` is not handling the recursivity well.
+        # Please see this Hypothesis issue:
+        # https://github.com/HypothesisWorks/hypothesis/issues/3026
+
 
 # noinspection PyUnusedLocal
 class TestRepresentationOfCondition(unittest.TestCase):