Merge pull request #360 from agriggio/smtlib-script-custom-logic

Allow the user to specify a custom logic for smtlibscript_from_formula()

pysmt/smtlib/script.py

@@ -22,10 +22,11 @@
 from six.moves import xrange
 
 import pysmt.smtlib.commands as smtcmd
-from pysmt.exceptions import UnknownSmtLibCommandError, NoLogicAvailableError
+from pysmt.exceptions import (UnknownSmtLibCommandError, NoLogicAvailableError,
+                              UndefinedLogicError)
 from pysmt.smtlib.printers import SmtPrinter, SmtDagPrinter, quote
 from pysmt.oracles import get_logic
-from pysmt.logics import get_closer_smtlib_logic
+from pysmt.logics import get_closer_smtlib_logic, Logic, SMTLIB2_LOGICS
 from pysmt.environment import get_env
 
 
@@ -228,21 +229,31 @@ def __str__(self):
         return "\n".join((str(cmd) for cmd in self.commands))
 
 
-def smtlibscript_from_formula(formula):
+def smtlibscript_from_formula(formula, logic=None):
     script = SmtLibScript()
 
-    # Get the simplest SmtLib logic that contains the formula
-    f_logic = get_logic(formula)
-
-    smt_logic = None
-    try:
-        smt_logic = get_closer_smtlib_logic(f_logic)
-    except NoLogicAvailableError:
-        warnings.warn("The logic %s is not reducible to any SMTLib2 " \
-                      "standard logic. Proceeding with non-standard " \
-                      "logic '%s'" % (f_logic, f_logic),
-                      stacklevel=3)
-        smt_logic = f_logic
+    if logic is None:
+        # Get the simplest SmtLib logic that contains the formula
+        f_logic = get_logic(formula)
+
+        smt_logic = None
+        try:
+            smt_logic = get_closer_smtlib_logic(f_logic)
+        except NoLogicAvailableError:
+            warnings.warn("The logic %s is not reducible to any SMTLib2 " \
+                          "standard logic. Proceeding with non-standard " \
+                          "logic '%s'" % (f_logic, f_logic),
+                          stacklevel=3)
+            smt_logic = f_logic
+    elif not (isinstance(logic, Logic) or isinstance(logic, str)):
+        raise UndefinedLogicError(str(logic))
+    else:
+        if logic not in SMTLIB2_LOGICS:
+            warnings.warn("The logic %s is not reducible to any SMTLib2 " \
+                          "standard logic. Proceeding with non-standard " \
+                          "logic '%s'" % (logic, logic),
+                          stacklevel=3)
+        smt_logic = logic
 
     script.add(name=smtcmd.SET_LOGIC,
                args=[smt_logic])

pysmt/test/smtlib/test_smtlibscript.py

@@ -25,6 +25,9 @@
 from pysmt.smtlib.script import smtlibscript_from_formula, evaluate_command
 from pysmt.smtlib.parser import get_formula_strict, get_formula, SmtLibParser
 from pysmt.solvers.smtlib import SmtLibIgnoreMixin
+from pysmt.logics import QF_UFLIRA
+from pysmt.exceptions import UndefinedLogicError
+
 
 
 class TestSmtLibScript(TestCase):
@@ -66,6 +69,24 @@ def test_from_formula(self):
         self.assertIn("(declare-fun y () Bool)", output)
         self.assertIn("(check-sat)", output)
 
+        # Use custom logic (as str)
+        script2 = smtlibscript_from_formula(f, logic="BOOL")
+        outstream = cStringIO()
+        script2.serialize(outstream)
+        output = outstream.getvalue()
+        self.assertIn("(set-logic BOOL)", output)
+
+        # Use custom logic (as Logic obj)
+        script3 = smtlibscript_from_formula(f, logic=QF_UFLIRA)
+        outstream = cStringIO()
+        script3.serialize(outstream)
+        output = outstream.getvalue()
+        self.assertIn("(set-logic QF_UFLIRA)", output)
+
+        # Custom logic must be a Logic or Str
+        with self.assertRaises(UndefinedLogicError):
+            smtlibscript_from_formula(f, logic=4)
+
 
     def test_get_strict_formula(self):