Prolog: allow custom inference limit in check_answers

1 files changed, 5 insertions, 4 deletions

diff --git a/prolog/engine.py b/prolog/engine.py
index e78bacb..c6ecab1 100644
--- a/prolog/engine.py
+++ b/prolog/engine.py
@@ -126,7 +126,7 @@ def pretty_vars(data):
 # specified by [answers] (a list of binding dictionaries) are returned. This
 # function succeeds if [query] finds each solution in [answers] at least once
 # within [timeout] seconds, and fails when it finds any other solution.
-def check_answers(engine, query, answers, timeout=10):
+def check_answers(engine, query, answers, timeout=10, inference_limit=100000):
     seen = []
     # Return false (only) if there is no expected answer where values of all
     # specified variables match values in [answer].
@@ -144,10 +144,11 @@ def check_answers(engine, query, answers, timeout=10):
     start = time.monotonic()
     try:
         # Limit inferences for each solution to curb unbounded recursion.
-        limited = '''(call_with_inference_limit(({}), 1000000, _ILR),
-                      _ILR \= inference_limit_exceeded)'''.format(query)
+        if inference_limit is not None:
+            query = '''(call_with_inference_limit(({}), {}, _ILR),
+                        _ILR \= inference_limit_exceeded)'''.format(query, inference_limit)
         # Run the query.
-        reply, output = ask(engine, limited, timeout)
+        reply, output = ask(engine, query, timeout)
         answer, error, more = process_answer(reply)
         if not check_answer(answer):
             return False