Added POW operator and generalized build_tree function

2 files changed, 46 insertions, 67 deletions

diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..485dee6
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1 @@
+.idea
diff --git a/tokenizer.py b/tokenizer.py
index 95af037..2a9a714 100644
--- a/tokenizer.py
+++ b/tokenizer.py
@@ -2,8 +2,18 @@ import sys
 
 
 STATE_ANY, STATE_OPERATOR, STATE_NUMBER, *_ = range(100)
-LEFT_PARENTHESIS, RIGHT_PARENTHESIS, NUMBER, MULTIPLY, PLUS, MINUS, *_ = range(100)
-OPERATORS = [MULTIPLY, PLUS, MINUS]
+(
+    LEFT_PARENTHESIS, RIGHT_PARENTHESIS, NUMBER,
+    MULTIPLY, PLUS, MINUS, POW,
+    *_
+) = range(100)
+OPERATORS = [MULTIPLY, PLUS, MINUS, POW]
+PRECEDENCE = {
+    PLUS: 10,
+    MINUS: 10,
+    MULTIPLY: 20,
+    POW: 30,
+}
 
 
 class Node:
@@ -30,6 +40,8 @@ class Node:
             return "Node(PLUS)"
         if self.type == MINUS:
             return "Node(MINUS)"
+        if self.type == POW:
+            return "Node(POW)"
     __str__ = __repr__
 
 
@@ -54,6 +66,8 @@ class Token:
             return "Token(PLUS)"
         if self.type == MINUS:
             return "Token(MINUS)"
+        if self.type == POW:
+            return "Token(POW)"
     __str__ = __repr__
 
 
@@ -95,15 +109,18 @@ def tokenize(line):
                 tokens.append(Token(MULTIPLY, None))
                 state = STATE_NUMBER
             else:
-                # print(1, tokens)
+                raise ValueError("Line is not a valid expression")
+        elif char == "^":
+            if state == STATE_OPERATOR:
+                tokens.append(Token(POW, None))
+                state = STATE_NUMBER
+            else:
                 raise ValueError("Line is not a valid expression")
         elif char in "-+0123456789":
             val, i = parse_number(line, i)
             tokens.append(Token(NUMBER, val))
             state = STATE_OPERATOR
         elif char != " ":
-            # print(char)
-            # print(2, tokens)
             raise ValueError("Line is not a valid expression")
         i += 1
 
@@ -111,26 +128,19 @@ def tokenize(line):
 
 
 def parse_parenthesis(tokens, start):
-    # print(tokens)
-    i = start
     end = start
     depth = 0
-    # print(tokens[i])
     while True:
-        token = tokens[i]
+        token = tokens[end]
         if token.type == RIGHT_PARENTHESIS:
             depth -= 1
         if token.type == LEFT_PARENTHESIS:
             depth += 1
         if depth == 0 and token.type == RIGHT_PARENTHESIS:
-            end = i + 1
             break
-        i += 1
-    # print("=== BUILDING PARENTHESIS ===")
-    # print(tokens[start + 1 : end - 1])
-    node = build_tree(tokens[start + 1 : end - 1])
-    # print("=== EXITED PARENTHESIS ===")
-    return node, end - 1
+        end += 1
+    node = build_tree(tokens[start + 1: end])
+    return node, end
 
 
 def build_tree(tokens):
@@ -138,74 +148,39 @@ def build_tree(tokens):
     current_node = None
     while i < len(tokens):
         token = tokens[i]
-        if token.type == LEFT_PARENTHESIS:
-            node, i = parse_parenthesis(tokens, i)
+        if token.type in [LEFT_PARENTHESIS, NUMBER]:
+            if token.type == LEFT_PARENTHESIS:
+                node, i = parse_parenthesis(tokens, i)
+            elif token.type == NUMBER:
+                node = Node(NUMBER, token.value)
+
             if current_node is None:
                 current_node = node
             elif current_node.type == NUMBER:
                 raise ValueError("Not a valid expression")
             elif current_node.type in OPERATORS:
                 if current_node.left is None:
-                    # print(f"Placed number {node} on the left of {current_node}")
                     current_node.left = node
                     node.parent = current_node
                 elif current_node.right is None:
-                    # print(f"Placed number {node} on the right of {current_node}")
                     current_node.right = node
                     node.parent = current_node
                     current_node = node
                 else:
                     raise ValueError("Not a valid expression")
-        elif token.type == NUMBER:
-            node = Node(NUMBER, token.value)
-            if current_node is None:
-                # print(f"Number {node} is new root")
-                current_node = node
-            elif current_node.type in OPERATORS:
-                if current_node.left is None:
-                    # print(f"Placed number {node} on the left of {current_node}")
-                    current_node.left = node
-                    node.parent = current_node
-                elif current_node.right is None:
-                    # print(f"Placed number {node} on the right of {current_node}")
-                    current_node.right = node
-                    node.parent = current_node
-                    current_node = node
-                else:
-                    raise ValueError("Not a valid expression")
-        elif token.type in [PLUS, MINUS]:
+
+        elif token.type in OPERATORS:
             node = Node(token.type, None)
             if current_node is None:
                 raise ValueError("Not a valid expression")
-            elif current_node.type == NUMBER:
+            elif current_node.type in [NUMBER] + OPERATORS:
                 while current_node.parent is not None:
+                    if PRECEDENCE[current_node.parent.type] < PRECEDENCE[node.type]:
+                        node.parent = current_node.parent
+                        if node.parent is not None:
+                            node.parent.right = node
+                        break
                     current_node = current_node.parent
-                # print(f"Operator {node} is new root with {current_node} on the left")
-                current_node.parent = node
-                node.left = current_node
-                current_node = node
-            elif current_node.type == MULTIPLY:
-                # print(f"Placed {node} over {node.left}")
-                current_node.parent = node
-                node.left = current_node
-                current_node = node
-            else:
-                raise ValueError("Not a valid expression")
-        elif token.type == MULTIPLY:
-            node = Node(MULTIPLY, None)
-            if current_node is None:
-                raise ValueError("Not a valid expression")
-            elif current_node.type in [NUMBER, MULTIPLY]:
-                # print(f"Replaced with {node} node {current_node}")
-                node.left = current_node
-                node.parent = current_node.parent
-                if node.parent is not None:
-                    node.parent.right = node
-                current_node.parent = node
-                current_node = node
-            elif current_node.type in [PLUS, MINUS]:
-                if current_node.parent is not None:
-                    raise ValueError("oops...")
                 current_node.parent = node
                 node.left = current_node
                 current_node = node
@@ -227,12 +202,15 @@ def expr_eval(node):
         return expr_eval(node.left) + expr_eval(node.right)
     elif node.type == MINUS:
         return expr_eval(node.left) - expr_eval(node.right)
+    elif node.type == POW:
+        return expr_eval(node.left) ** expr_eval(node.right)
+
 
 def main(*argv):
     data = argv[0]
     tokens = tokenize(data)
-    node = build_tree(tokens)
-    print(expr_eval(node))
+    node1 = build_tree(tokens)
+    print(expr_eval(node1))
 
 
 if __name__ == "__main__":