Explorar o código

ExpressionBase magic methods convert the other arguments to expressionnodes.

Sander Mathijs van Veen %!s(int64=14) %!d(string=hai) anos
pai
achega
0065eb9eda
Modificáronse 1 ficheiros con 12 adicións e 5 borrados
  1. 12 5
      src/node.py

+ 12 - 5
src/node.py

@@ -52,6 +52,10 @@ OP_MAP = {
         }
 
 
+def to_expression(obj):
+    return obj if isinstance(obj, ExpressionBase) else ExpressionLeaf(obj)
+
+
 class ExpressionBase(object):
     def __lt__(self, other):
         """
@@ -116,19 +120,22 @@ class ExpressionBase(object):
         return self.is_leaf() and self.type & (TYPE_FLOAT | TYPE_INTEGER)
 
     def __add__(self, other):
-        return ExpressionNode('+', self, other)
+        return ExpressionNode('+', self, to_expression(other))
 
     def __sub__(self, other):
-        return ExpressionNode('-', self, other)
+        return ExpressionNode('-', self, to_expression(other))
 
     def __mul__(self, other):
-        return ExpressionNode('*', self, other)
+        return ExpressionNode('*', self, to_expression(other))
 
     def __div__(self, other):
-        return ExpressionNode('-', self, other)
+        return ExpressionNode('-', self, to_expression(other))
 
     def __pow__(self, other):
-        return ExpressionNode('^', self, other)
+        return ExpressionNode('^', self, to_expression(other))
+
+    def __neg__(self):
+        return ExpressionNode('-', to_expression(self))
 
 
 class ExpressionNode(Node, ExpressionBase):