Added negation checks to Node.

src/node.py

@@ -104,9 +104,27 @@ class ExpressionBase(object):
     def is_op(self, op):
         return not self.is_leaf() and self.op == op
 
+    def is_op_or_negated(self, op):
+        if self.is_leaf():
+            return False
+
+        if self.op == OP_NEG:
+            return self[0].is_op(op)
+
+        return self.op == op
+
     def is_leaf(self):
         return self.type != TYPE_OPERATOR
 
+    def is_leaf_or_negated(self):
+        if  self.is_leaf():
+            return True
+
+        if self.is_op(OP_NEG):
+            return self[0].is_leaf()
+
+        return False
+
     def is_power(self):
         return not self.is_leaf() and self.op == OP_POW
 

tests/test_node.py

 import unittest
 
-from src.node import ExpressionNode as N, ExpressionLeaf as L
+from src.node import ExpressionNode as N, ExpressionLeaf as L, OP_ADD
 from tests.rulestestcase import tree
 
 
@@ -23,7 +23,26 @@ class TestNode(unittest.TestCase):
                         < N('*', L(3), N('^', L('a'), L('b'))))
         self.assertFalse(N('^', L('a'), L(3)) < N('^', L('a'), L(2)))
 
-    def test_is_power_true(self):
+    def test_is_op(self):
+        self.assertTrue(N('+', *self.l[:2]).is_op(OP_ADD))
+        self.assertFalse(N('-', *self.l[:2]).is_op(OP_ADD))
+
+    def test_is_op_or_negated(self):
+        self.assertTrue(N('+', *self.l[:2]).is_op_or_negated(OP_ADD))
+        self.assertTrue(N('-', N('+', *self.l[:2])).is_op_or_negated(OP_ADD))
+        self.assertFalse(N('-', *self.l[:2]).is_op_or_negated(OP_ADD))
+
+    def test_is_leaf(self):
+        self.assertTrue(L(2).is_leaf())
+        self.assertFalse(N('+', *self.l[:2]).is_leaf())
+
+    def test_is_leaf_or_negated(self):
+        self.assertTrue(L(2).is_leaf_or_negated())
+        self.assertTrue(N('-', L(2)).is_leaf_or_negated())
+        self.assertFalse(N('+', *self.l[:2]).is_leaf_or_negated())
+        self.assertFalse(N('-', N('+', *self.l[:2])).is_leaf_or_negated())
+
+    def test_is_power(self):
         self.assertTrue(N('^', *self.l[:2]).is_power())
         self.assertFalse(N('+', *self.l[:2]).is_power())