fix codegen, so it can handle operators

1 files changed, 17 insertions, 5 deletions

diff --git a/util/codegen.py b/util/codegen.py
index 3f307f2..bd69202 100644
--- a/util/codegen.py
+++ b/util/codegen.py
@@ -30,6 +30,8 @@
 
     :copyright: Copyright 2008 by Armin Ronacher.
     :license: BSD.
+
+    Modified by Marion Zepf.
 """
 from ast import *
 
@@ -63,6 +65,16 @@ class SourceGenerator(NodeVisitor):
     `node_to_source` function.
     """
 
+    UNARYOP_SYMBOLS = {Invert: "~", Not: "not", UAdd: "+", USub: "-"}
+    # TODO avoid turning (-1)**2 into -1**2
+    BINOP_SYMBOLS = {Add: "+", Sub: "-", Mult: "*", Div: "/", Mod: "%",
+                     LShift: "<<", RShift:">>", BitOr: "|", BitXor: "^",
+                     BitAnd: "&", FloorDiv: "//", Pow: "**"}
+    BOOLOP_SYMBOLS = {And: "and", Or: "or"}
+    CMPOP_SYMBOLS = {Eq: "==", NotEq: "!=", Lt: "<", LtE: "<=", Gt: ">",
+                     GtE: ">=", Is: "is", IsNot: "is not", In: "in",
+                     NotIn: "not in"}
+
     def __init__(self, indent_with, add_line_information=False):
         self.result = []
         self.indent_with = indent_with
@@ -142,7 +154,7 @@ class SourceGenerator(NodeVisitor):
     def visit_AugAssign(self, node):
         self.newline(node)
         self.visit(node.target)
-        self.write(BINOP_SYMBOLS[type(node.op)] + '=')
+        self.write(self.BINOP_SYMBOLS[node.op] + '=')
         self.visit(node.value)
 
     def visit_ImportFrom(self, node):
@@ -422,14 +434,14 @@ class SourceGenerator(NodeVisitor):
 
     def visit_BinOp(self, node):
         self.visit(node.left)
-        self.write(' %s ' % BINOP_SYMBOLS[type(node.op)])
+        self.write(' %s ' % self.BINOP_SYMBOLS[type(node.op)])
         self.visit(node.right)
 
     def visit_BoolOp(self, node):
         self.write('(')
         for idx, value in enumerate(node.values):
             if idx:
-                self.write(' %s ' % BOOLOP_SYMBOLS[type(node.op)])
+                self.write(' %s ' % self.BOOLOP_SYMBOLS[node.op])
             self.visit(value)
         self.write(')')
 
@@ -437,13 +449,13 @@ class SourceGenerator(NodeVisitor):
         self.write('(')
         self.write(node.left)
         for op, right in zip(node.ops, node.comparators):
-            self.write(' %s %%' % CMPOP_SYMBOLS[type(op)])
+            self.write(' %s %%' % self.CMPOP_SYMBOLS[op])
             self.visit(right)
         self.write(')')
 
     def visit_UnaryOp(self, node):
         self.write('(')
-        op = UNARYOP_SYMBOLS[type(node.op)]
+        op = self.UNARYOP_SYMBOLS[node.op]
         self.write(op)
         if op == 'not':
             self.write(' ')