Improve handling of Prolog exceptions

1 files changed, 58 insertions, 45 deletions

diff --git a/prolog/engine.py b/prolog/engine.py
index 26215f8..7c90d1a 100644
--- a/prolog/engine.py
+++ b/prolog/engine.py
@@ -112,74 +112,87 @@ class PrologEngine(object):
         # Dictionary of correct answers; keys are (pid, query).
         self.answers = {}
 
-    def exception(self, qid):
-        ref = PL_exception(qid)
-        if not ref:
-            return False
-        try:
-            ex = Term(ref=ref)
-            msg = Term()
-            if PL_call_predicate(None, self.err_flags,
-                                 self.p['message_to_string/2'], Termv([ex, msg]).ref):
-                raise Exception(str(msg))
-            raise Exception('Unknown error')
-        finally:
-            PL_clear_exception()
-        return False
+    # Return a description of the last exception, or None if no error occurred.
+    def error(self, qid):
+        error_ref = PL_exception(qid)
+        if not error_ref:
+            return None
+        PL_clear_exception()
+
+        # Get the Prolog error message.
+        fid = PL_open_foreign_frame()
+        msg = Term()
+        if PL_call_predicate(None, self.err_flags, self.p['message_to_string/2'],
+                             Termv([Term(ref=error_ref), msg]).ref):
+            error_str = str(msg)
+        else:
+            error_str = 'Unknown error'
+        PL_discard_foreign_frame(fid)
+
+        return error_str
 
+    # Call the Prolog predicate [name]. Raise an exception on error.
     def call(self, name, args):
         qid = PL_open_query(None, self.err_flags, self.p[name], Termv(args).ref)
         try:
             if not PL_next_solution(qid):
-                if not self.exception(qid):
-                    return False
+                error_msg = self.error(qid)
+                if error_msg:
+                    raise Exception(error_msg)
+                return False
         finally:
             PL_cut_query(qid)
         return True
 
-    # Get up to [n] solutions to query [q].
+    # Get up to [n] solutions to query [q]. If there are no solutions, return
+    # an empty list. Raise an exception on error (either from self.call, or due
+    # to malformed/unsafe query or a timeout).
     def query(self, q, module=None, n=1):
         if module is not None:
             q = '{}:({})'.format(module, q)
 
         fid = PL_open_foreign_frame()
-
-        var_names = Term()
-        options = Term([Term('variable_names', [var_names])])
-        goal = Term()
+        qid = None
         try:
-            # Parse term and store variable names.
+            # Parse the query and store variable names.
+            goal = Term()
+            var_names = Term()
+            options = Term([Term('variable_names', [var_names])])
             if not self.call('read_term_from_atom/3', [Term(Atom(q)), goal, options]):
                 raise Exception('Warning: Could not read term from {}\n'.format(q))
+
             # Check if goal is safe with currently loaded rules.
             if not self.call('safe_goal/1', [goal]):
                 raise Exception('Warning: Unsafe goal: {}\n'.format(goal))
-        except Exception as ex:
-            raise
 
-        solutions = Term()
-        goal_aux = Term('findnsols', [Term(n), goal, goal, solutions])
-        qid = PL_open_query(None, self.err_flags, self.p['call_with_time_limit/2'],
-                            Termv([Term(0.01), goal_aux]).ref)
+            solutions = Term()
+            goal_aux = Term('findnsols', [Term(n), goal, goal, solutions])
+            qid = PL_open_query(None, self.err_flags, self.p['call_with_time_limit/2'],
+                                Termv([Term(0.01), goal_aux]).ref)
 
-        result = []
-        if PL_next_solution(qid):
-            for solution in solutions:
+            result = []
+            if PL_next_solution(qid):
                 fid_solution = PL_open_foreign_frame()
-                PL_unify(goal.ref, solution.ref)
-                variables = {}
-                for var in var_names:
-                    name, value = Term(), Term()
-                    PL_get_arg(1, var.ref, name.ref)
-                    PL_get_arg(2, var.ref, value.ref)
-                    variables[str(name)] = str(value)
-                result.append(variables)
-                PL_rewind_foreign_frame(fid_solution)
-        else:
-            self.exception(qid)
-
-        PL_close_query(qid)
-        PL_discard_foreign_frame(fid)
+                for solution in solutions:
+                    PL_unify(goal.ref, solution.ref)
+                    variables = {}
+                    for var in var_names:
+                        name, value = Term(), Term()
+                        PL_get_arg(1, var.ref, name.ref)
+                        PL_get_arg(2, var.ref, value.ref)
+                        variables[str(name)] = str(value)
+                    result.append(variables)
+                    PL_rewind_foreign_frame(fid_solution)
+                PL_discard_foreign_frame(fid_solution)
+            else:
+                # Check for exceptions.
+                error_msg = self.error(qid)
+                if error_msg:
+                    raise Exception(error_msg)
+        finally:
+            if qid:
+                PL_close_query(qid)
+            PL_discard_foreign_frame(fid)
 
         return result