Add prolog.engine.run

Collects all the answers to a given query within the specified timeout.

1 files changed, 85 insertions, 4 deletions

diff --git a/prolog/engine.py b/prolog/engine.py
index f8e563f..b0dc1aa 100644
--- a/prolog/engine.py
+++ b/prolog/engine.py
@@ -4,7 +4,10 @@ import collections
 import html
 import http.client
 import json
+from operator import itemgetter
 import re
+import socket
+import time
 import urllib
 
 def strip_html(text):
@@ -123,9 +126,87 @@ def pretty_vars(data):
         result += [strip_html(b) for b in data['residuals']]
     return ',\n'.join(result) if result else 'true'
 
+# Get all solutions to [query] given background knowledge [code] that are found
+# within [timeout] seconds.
+def run(code, query, timeout):
+    # Returns a tuple ((bindings, constraints), error, more?) for one answer.
+    def process_answer(reply):
+        # When an engine is destroyed, a nested data object has the actual
+        # query result.
+        if reply['event'] == 'destroy':
+            reply = reply['data']
+
+        if reply['event'] == 'success':
+            # Return a dictionary {var: value} and a list of constraints (as
+            # strings) from the JSON object returned by Prolog.
+            data = reply['data'][0]
+            bindings = {}
+            for binding in data['variables']:
+                value = strip_html(binding['value'])
+                for name in binding['variables']:
+                    bindings[name] = value
+            constraints = [strip_html(r) for r in data.get('residuals', [])]
+            return (bindings, constraints), None, reply['more']
+        elif reply['event'] == 'failure':
+            return None, None, False
+        elif reply['event'] == 'error':
+            # Remove potential module name (engine ID) from the error message.
+            error = ('error', reply['data'].replace("'{}':".format(reply['id']), ''))
+            return None, error, False
+
+    start = time.monotonic()
+    result, messages = [], []
+    engine = None
+    try:
+        # Create a new pengine.
+        reply, output = create(code=code, timeout=timeout)
+        messages += output
+        if reply.get('event') != 'create':
+            raise Exception('System error: creating pengine')
+        engine = reply['id']
+        if 'error' in map(itemgetter(0), messages):
+            return None, messages
+
+        # Run the query.
+        real_timeout = timeout - (time.monotonic()-start)
+        if real_timeout <= 0:
+            raise socket.timeout()
+        reply, output = ask(engine, query, real_timeout)
+        messages += output
+        if 'error' in map(itemgetter(0), output):
+            return None, messages
+
+        bindings, error, more = process_answer(reply)
+        if bindings:
+            result.append(bindings)
+        if error:
+            messages.append(error)
+
+        # Continue while there are more potential answers.
+        while more:
+            real_timeout = timeout - (time.monotonic()-start)
+            if real_timeout <= 0:
+                raise socket.timeout()
+            reply, output = next(engine, timeout=real_timeout)
+            messages += output
+            bindings, error, more = process_answer(reply)
+            if bindings:
+                result.append(bindings)
+            if error:
+                messages.append(error)
+    except socket.timeout as ex:
+        result.append('timeout')
+        abort(engine)
+    finally:
+        if engine:
+            destroy(engine)
+
+    return result, messages
+
+
 # Basic sanity check.
 if __name__ == '__main__':
-    engine = create(code='dup([],[]). dup([H|T],[H,H|TT]) :- dup(T,TT).')['id']
-    print('engine id is ' + engine)
-    print(ask(engine, "run_tests({}, Result)".format('dup/2')))
-    destroy(engine)
+    answers, messages = run('b(Y). a(X) :- {X > 3}, (X = 5 ; {X > 4}).', 'a(X)', timeout=1)
+    print(messages)
+    for bindings, constraints in answers:
+        print('bindings: {}, constraints: {}'.format(bindings, constraints))