Repository: pysmt/pysmt
Pull request #108 (2015-05-08 14:31:15)

FNODE: Extended support to infix notation as per #92.
Added support for the follwing infix operators:
* ```x & y```
* ```x | y```
* ```~x```
* ```-x```
* ```int(x)```
* ```float(x)```
* ```long(x)```

And added support for automatic propotion of python constants to pysmt constants:
E.g.
```python
1.5 + Symbol(x, Real) <= 8
```


Patch:
diff --git a/pysmt/fnode.py b/pysmt/fnode.py
index a6f7044..66cb406 100644
--- a/pysmt/fnode.py
+++ b/pysmt/fnode.py
@@ -37,6 +37,7 @@ from pysmt.operators import (FORALL, EXISTS, AND, OR, NOT, IMPLIES, IFF,
                              BV_OPERATORS)
 from pysmt.typing import BOOL, REAL, INT, PYSMT_TYPES, BVType
 from pysmt.decorators import deprecated
+from pysmt.utils import is_python_integer, is_python_rational, is_python_boolean
 
 FNodeContent = collections.namedtuple("FNodeContent",
                                       ["node_type", "args", "payload"])
@@ -352,7 +353,19 @@ class FNode(object):
 
     # Infix Notation
     def _apply_infix(self, right, function):
-        if pysmt.shortcuts.get_env().enable_infix_notation:
+        env = pysmt.shortcuts.get_env()
+        if env.enable_infix_notation:
+            if is_python_boolean(right):
+                right = pysmt.shortcuts.Bool(right)
+            elif is_python_integer(right):
+                ty = env.stc.get_type(self)
+                if ty.is_real_type():
+                    right = pysmt.shortcuts.Real(right)
+                else:
+                    right = pysmt.shortcuts.Int(right)
+            elif is_python_rational(right):
+                right = pysmt.shortcuts.Real(right)
+
             return function(self, right)
         else:
             raise Exception("Cannot use infix notation")
@@ -378,12 +391,22 @@ class FNode(object):
     def __add__(self, right):
         return self._apply_infix(right, pysmt.shortcuts.Plus)
 
+    def __radd__(self, right):
+        return self._apply_infix(right, pysmt.shortcuts.Plus)
+
     def __sub__(self, right):
         return self._apply_infix(right, pysmt.shortcuts.Minus)
 
+    def __rsub__(self, right):
+        minus_self = -self
+        return minus_self._apply_infix(right, pysmt.shortcuts.Plus)
+
     def __mul__(self, right):
         return self._apply_infix(right, pysmt.shortcuts.Times)
 
+    def __rmul__(self, right):
+        return self._apply_infix(right, pysmt.shortcuts.Times)
+
     def __div__(self, right):
         return self._apply_infix(right, pysmt.shortcuts.Div)
 
@@ -402,4 +425,45 @@ class FNode(object):
     def __le__(self, right):
         return self._apply_infix(right, pysmt.shortcuts.LE)
 
+    def __and__(self, other):
+        return self.And(other)
+
+    def __rand__(self, other):
+        return self.And(other)
+
+    def __or__(self, other):
+        return self.Or(other)
+
+    def __ror__(self, other):
+        return self.Or(other)
+
+    def __xor__(self, other):
+        return self._apply_infix(other, pysmt.shortcuts.Xor)
+
+    def __neg__(self):
+        return self._apply_infix(-1, pysmt.shortcuts.Times)
+
+    def __invert__(self):
+        env = pysmt.shortcuts.get_env()
+        if not env.enable_infix_notation:
+            raise Exception("Cannot use infix notation")
+        return pysmt.shortcuts.Not(self)
+
+    def __int__(self):
+        if self.is_int_constant():
+            return self.constant_value()
+        raise NotImplementedError("Cannot convert `%s` to integer" % str(self))
+
+    def __long__(self):
+        if self.is_int_constant():
+            return self.constant_value()
+        raise NotImplementedError("Cannot convert `%s` to integer" % str(self))
+
+    def __float__(self):
+        if self.is_int_constant() or self.is_real_constant():
+            return float(self.constant_value())
+        raise NotImplementedError("Cannot convert `%s` to float" % str(self))
+
+
+
 # EOC FNode
diff --git a/pysmt/formula.py b/pysmt/formula.py
index 2f3a6ad..9d10651 100644
--- a/pysmt/formula.py
+++ b/pysmt/formula.py
@@ -34,7 +34,7 @@ from fractions import Fraction
 import pysmt.typing as types
 import pysmt.operators as op
 
-from pysmt.utils import is_integer
+from pysmt.utils import is_python_integer
 from pysmt.fnode import FNode, FNodeContent
 from pysmt.exceptions import NonLinearError
 from pysmt.walkers.identitydag import IdentityDagWalker
@@ -315,7 +315,7 @@ class FormulaManager(object):
             val = value
         elif type(value) == tuple:
             val = Fraction(value[0], value[1])
-        elif is_integer(value):
+        elif is_python_integer(value):
             val = Fraction(value, 1)
         elif type(value) == float:
             val = Fraction.from_float(value)
@@ -334,7 +334,7 @@ class FormulaManager(object):
         if value in self.int_constants:
             return self.int_constants[value]
 
-        if is_integer(value):
+        if is_python_integer(value):
             n = self.create_node(node_type=op.INT_CONSTANT,
                                  args=tuple(),
                                  payload=value)
@@ -532,7 +532,7 @@ class FormulaManager(object):
         if width is None:
             raise ValueError("Need to specify a width for the constant")
 
-        if is_integer(value):
+        if is_python_integer(value):
             if value < 0:
                 raise ValueError("Cannot specify a negative value: %d" % value)
             if value >= 2**width:
@@ -587,7 +587,7 @@ class FormulaManager(object):
     def BVExtract(self, formula, start=0, end=None):
         """Returns the slice of formula from start to end (inclusive)."""
         if end is None: end = formula.bv_width()-1
-        assert is_integer(start) and is_integer(end)
+        assert is_python_integer(start) and is_python_integer(end)
         assert end >= start and start >= 0, "Start: %d ; End: %d" % (start,end)
         size = end-start+1
 
@@ -655,7 +655,7 @@ class FormulaManager(object):
 
     def BVLShl(self, left, right):
         """Returns the logical left shift the BV."""
-        if is_integer(right):
+        if is_python_integer(right):
             right = self.BV(right, left.bv_width())
         return self.create_node(node_type=op.BV_LSHL,
                                 args=(left, right),
@@ -663,7 +663,7 @@ class FormulaManager(object):
 
     def BVLShr(self, left, right):
         """Returns the logical right shift the BV."""
-        if is_integer(right):
+        if is_python_integer(right):
             right = self.BV(right, left.bv_width())
         return self.create_node(node_type=op.BV_LSHR,
                                 args=(left, right),
@@ -671,7 +671,7 @@ class FormulaManager(object):
 
     def BVRol(self, formula, steps):
         """Returns the LEFT rotation of the BV by the number of steps."""
-        if not is_integer(steps):
+        if not is_python_integer(steps):
             raise TypeError("BVRol: 'steps' should be an integer. Got %s" % steps)
         return self.create_node(node_type=op.BV_ROL,
                                 args=(formula,),
@@ -679,7 +679,7 @@ class FormulaManager(object):
 
     def BVRor(self, formula, steps):
         """Returns the RIGHT rotation of the BV by the number of steps."""
-        if not is_integer(steps):
+        if not is_python_integer(steps):
             raise TypeError("BVRor: 'steps' should be an integer. Got %s" % steps)
         return self.create_node(node_type=op.BV_ROR,
                                 args=(formula,),
@@ -690,7 +690,7 @@ class FormulaManager(object):
 
         New bits are set to zero.
         """
-        if not is_integer(increase):
+        if not is_python_integer(increase):
             raise TypeError("BVZext: 'increase' should be an integer. Got %s" % increase)
         return self.create_node(node_type=op.BV_ZEXT,
                                 args=(formula,),
@@ -702,7 +702,7 @@ class FormulaManager(object):
 
         New bits are set according to the most-significant-bit.
         """
-        if not is_integer(increase):
+        if not is_python_integer(increase):
             raise TypeError("BVSext: 'increase' should be an integer. Got %s" % increase)
         return self.create_node(node_type=op.BV_SEXT,
                                 args=(formula,),
diff --git a/pysmt/utils.py b/pysmt/utils.py
index c738523..5e0223a 100644
--- a/pysmt/utils.py
+++ b/pysmt/utils.py
@@ -1,13 +1,22 @@
 import io
 from six import PY2
+from fractions import Fraction
 
-def is_integer(n):
+def is_python_integer(n):
     if PY2:
         return type(n) == long or type(n) == int
     else:
         return type(n) == int
 
 
+def is_python_rational(n):
+    return is_python_integer(n) or type(n) == float or type(n) == Fraction
+
+
+def is_python_boolean(n):
+    return n is True or n is False
+
+
 def BufferedTextReader(fname):
     if PY2:
         return io.BufferedReader(io.FileIO(fname, 'r'))


Test patch:
diff --git a/pysmt/test/test_formula.py b/pysmt/test/test_formula.py
index c9106c4..6616826 100644
--- a/pysmt/test/test_formula.py
+++ b/pysmt/test/test_formula.py
@@ -21,7 +21,8 @@ from six.moves import xrange
 
 import pysmt
 from pysmt.typing import BOOL, REAL, INT, FunctionType
-from pysmt.shortcuts import Symbol, is_sat, is_valid, Implies, GT, Plus
+from pysmt.shortcuts import Symbol, is_sat, Not, Implies, GT, Plus, Int, Real
+from pysmt.shortcuts import Minus, Times, Xor, And, Or, TRUE
 from pysmt.shortcuts import get_env
 from pysmt.environment import Environment
 from pysmt.test import TestCase, skipIfNoSolverForLogic
@@ -623,6 +624,46 @@ class TestFormulaManager(TestCase):
 
         self.assertEqual(p + p, Plus(p,p))
         self.assertEqual(p > p, GT(p,p))
+        get_env().enable_infix_notation = False
+
+    def test_infix_extended(self):
+        p, r, x, y = self.p, self.r, self.x, self.y
+        get_env().enable_infix_notation = True
+
+        self.assertEqual(Plus(p, Int(1)), p + 1)
+        self.assertEqual(Plus(r, Real(1)), r + 1)
+        self.assertEqual(Times(r, Real(1)), r * 1)
+
+        self.assertEqual(Minus(p, Int(1)), p - 1)
+        self.assertEqual(Minus(r, Real(1)), r - 1)
+        self.assertEqual(Times(r, Real(1)), r * 1)
+
+        self.assertEqual(Plus(r, Real(1.5)), r + 1.5)
+        self.assertEqual(Minus(r, Real(1.5)), r - 1.5)
+        self.assertEqual(Times(r, Real(1.5)), r * 1.5)
+
+        self.assertEqual(Plus(r, Real(1.5)), 1.5 + r)
+        self.assertEqual(Times(r, Real(1.5)), 1.5 * r)
+
+        with self.assertRaises(TypeError):
+            foo = p + 1.5
+
+        self.assertEqual(Not(x), ~x)
+
+        self.assertEqual(Times(r, Real(-1)), -r)
+        self.assertEqual(Times(p, Int(-1)), -p)
+
+        self.assertEqual(Xor(x, y), x ^ y)
+        self.assertEqual(And(x, y), x & y)
+        self.assertEqual(Or(x, y), x | y)
+
+        self.assertEqual(Or(x, TRUE()), x | True)
+        self.assertEqual(Or(x, TRUE()), True | x)
+
+        self.assertEqual(And(x, TRUE()), x & True)
+        self.assertEqual(And(x, TRUE()), True & x)
+
+        get_env().enable_infix_notation = False
 
     def test_toReal(self):
         f = self.mgr.Equals(self.rconst, self.mgr.ToReal(self.p))