support multiple ir_types in ir_nodes

HypothesisWorks · Jan 8, 2025 · 4348fa7 · 4348fa7
1 parent 85c823a
commit 4348fa7
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/hypothesis-python/tests/conjecture/common.py b/hypothesis-python/tests/conjecture/common.py
@@ -356,10 +356,11 @@ def draw_value(ir_type, kwargs):
 
 
 @st.composite
-def ir_nodes(draw, *, was_forced=None, ir_type=None):
-    if ir_type is None:
+def ir_nodes(draw, *, was_forced=None, ir_types=None):
+    if ir_types is None:
         (ir_type, kwargs) = draw(ir_types_and_kwargs())
     else:
+        ir_type = draw(st.sampled_from(ir_types))
         kwargs = draw(kwargs_strategy(ir_type))
     # ir nodes don't include forced in their kwargs. see was_forced attribute
     del kwargs["forced"]

diff --git a/hypothesis-python/tests/conjecture/test_data_tree.py b/hypothesis-python/tests/conjecture/test_data_tree.py
@@ -579,7 +579,7 @@ def _draw(data, node, *, forced=None):
     return getattr(data, f"draw_{node.ir_type}")(**node.kwargs, forced=forced)
 
 
-@given(ir_nodes(was_forced=True, ir_type="float"))
+@given(ir_nodes(was_forced=True, ir_types=["float"]))
 def test_simulate_forced_floats(node):
     tree = DataTree()