PrologEngine: add load/unload functions

1 files changed, 37 insertions, 29 deletions

diff --git a/prolog/engine.py b/prolog/engine.py
index 85c5559..597ebc4 100644
--- a/prolog/engine.py
+++ b/prolog/engine.py
@@ -135,29 +135,23 @@ class PrologEngine(object):
         self.problems[pid] = Problem(name, solution, facts, tests)
 
         # Load the solution in 'solution<pid>' module.
-        m_problem = 'problem{}'.format(pid)
+        mod_problem = 'problem{}'.format(pid)
 
         fid = PL_open_foreign_frame()
-        predicates = set()
-        for rule in prolog.util.split(solution):
-            self.call('assertz/1', [Term('{}:({})'.format(m_problem, rule))])
-            predicates.add('{}:{}'.format(m_problem, self.predicate_indicator(rule)))
-        if facts:
-            for rule in prolog.util.split(facts):
-                if rule not in self.facts:
-                    self.call('assertz/1', [Term(rule)])
-                    predicates.add(self.predicate_indicator(rule))
-                    self.facts.add(rule)
+        predicates = self.load(solution, mod_problem)
+        if facts and facts not in self.facts:
+            predicates |= self.load(facts)
+            self.facts.add(facts)
         self.call('compile_predicates/1', [Term([Term(p) for p in predicates])])
 
         # Import solutions for dependency predicates.
         for i in depends:
-            m_dependency = 'problem{}'.format(i)
-            self.call('add_import_module/3', [Term(m_problem), Term(m_dependency), Term('end')])
+            mod_dependency = 'problem{}'.format(i)
+            self.call('add_import_module/3', [Term(mod_problem), Term(mod_dependency), Term('end')])
 
         # Find the correct test answers.
         for query in tests:
-            result = self.query(query, m_problem)
+            result = self.query(query, mod_problem)
             if result is None or len(result) < 1 or 'X' not in result[0]:
                 raise Exception('Error finding correct answer to query "{}"'.format(query))
             self.problems[pid].tests[query] = result[0]['X']
@@ -165,11 +159,11 @@ class PrologEngine(object):
 
     # Import the correct solution for problem [pid] into module for user [uid].
     def mark_solved(self, uid, pid):
-        m_user = 'user{}'.format(uid)
-        m_problem = 'problem{}'.format(pid)
+        mod_user = 'user{}'.format(uid)
+        mod_problem = 'problem{}'.format(pid)
 
         fid = PL_open_foreign_frame()
-        result = self.call('add_import_module/3', [Term(m_user), Term(m_problem), Term('end')])
+        result = self.call('add_import_module/3', [Term(mod_user), Term(mod_problem), Term('end')])
         PL_discard_foreign_frame(fid)
         return result
 
@@ -229,20 +223,15 @@ class PrologEngine(object):
     # Test whether [code] gives the same answer to [query] as the correct
     # solution for problem [pid]. The solution should be loaded beforehand.
     def test(self, uid, pid, code):
-        # Module name for user code.
-        m_user = 'user{}'.format(uid)
+        mod_user = 'user{}'.format(uid)
 
         fid = PL_open_foreign_frame()
         correct = True
         predicates = set()
         try:
-            # Load the user program [code].
-            for rule in prolog.util.split(code):
-                self.call('assertz/1', [Term('{}:({})'.format(m_user, rule))])
-                predicates.add(self.predicate_indicator(rule))
-
+            self.load(code, mod_user, predicates)
             for query, answer in sorted(self.problems[pid].tests.items()):
-                result = self.query(query, m_user, n=1)
+                result = self.query(query, mod_user, n=1)
                 if len(result) != 1 or result[0]['X'] != answer:
                     correct = False
                     break
@@ -250,7 +239,7 @@ class PrologEngine(object):
                 # If a correct solution was found, see if another (incorrect)
                 # solution is found in the first 10 answers.
                 try:
-                    result = self.query(query, m_user, n=10)
+                    result = self.query(query, mod_user, n=10)
                     unique = set([r['X'] for r in result])
                     if len(unique) != 1:
                         correct = False
@@ -262,9 +251,7 @@ class PrologEngine(object):
         except Exception as ex:
             correct = False
 
-        # Unload all loaded rules.
-        for predicate in predicates:
-            self.call('abolish/1', [Term(m_user + ':' + predicate)])
+        self.unload(predicates)
         PL_discard_foreign_frame(fid)
 
         return correct
@@ -283,6 +270,27 @@ class PrologEngine(object):
             PL_cut_query(qid)
         return True
 
+    # Load rules from [program] into [module] and return the corresponding
+    # predicate names. Since this function might not return due to exception,
+    # the [predicates] argument can be passed where the names will be stored.
+    def load(self, program, module=None, predicates=None):
+        if predicates is None:
+            predicates = set()
+        for rule in prolog.util.split(program):
+            name = self.predicate_indicator(rule)
+            if module:
+                rule = '{}:({})'.format(module, rule)
+                name = '{}:{}'.format(module, name)
+            self.call('assertz/1', [Term(rule)])
+            predicates.add(name)
+        return predicates
+
+    # Unload and remove the "dynamic" property for all rules implementing
+    # [predicates].
+    def unload(self, predicates):
+        for predicate in predicates:
+            self.call('abolish/1', [Term(predicate)])
+
     # Return a description of the last exception, or None if no error occurred.
     def error(self, qid):
         error_ref = PL_exception(qid)