m-labs · Oct 14, 2014 · Oct 14, 2014
Showing with 37 additions and 3 deletions.

+31 −3 artiq/transforms/fold_constants.py

+6 −0 artiq/transforms/tools.py
diff --git a/artiq/transforms/fold_constants.py b/artiq/transforms/fold_constants.py
@@ -29,6 +29,12 @@
 }
 
 
+_ast_boolops = {
+    ast.Or: lambda x, y: x or y,
+    ast.And: lambda x, y: x and y
+}
+
+
 class _ConstantFolder(ast.NodeTransformer):
     def visit_UnaryOp(self, node):
         self.generic_visit(node)
@@ -62,13 +68,35 @@ def visit_BinOp(self, node):
             return node
         return ast.copy_location(result, node)
 
+    def visit_BoolOp(self, node):
+        self.generic_visit(node)
+        new_values = []
+        for value in node.values:
+            try:
+                value_c = eval_constant(value)
+            except NotConstant:
+                new_values.append(value)
+            else:
+                if new_values and not isinstance(new_values[-1], ast.AST):
+                    op = _ast_boolops[type(node.op)]
+                    new_values[-1] = op(new_values[-1], value_c)
+                else:
+                    new_values.append(value_c)
+        new_values = [v if isinstance(v, ast.AST) else value_to_ast(v)
+                      for v in new_values]
+        if len(new_values) > 1:
+            node.values = new_values
+            return node
+        else:
+            return new_values[0]
+
     def visit_Call(self, node):
         self.generic_visit(node)
         fn = node.func.id
         constant_ops = {
-            "int":    int,
-            "int64":  int64,
-            "round":  round,
+            "int": int,
+            "int64": int64,
+            "round": round,
             "round64": round64
         }
         if fn in constant_ops:

diff --git a/artiq/transforms/tools.py b/artiq/transforms/tools.py
@@ -19,6 +19,10 @@ def value_to_ast(value):
             func=ast.Name("int64", ast.Load()),
             args=[ast.Num(int(value))],
             keywords=[], starargs=None, kwargs=None)
+    elif isinstance(value, bool) or value is None:
+        # must also be before int
+        # isinstance(True/False, int) == True
+        return ast.NameConstant(value)
     elif isinstance(value, (int, float)):
         return ast.Num(value)
     elif isinstance(value, Fraction):
@@ -49,6 +53,8 @@ def eval_constant(node):
         return node.n
     elif isinstance(node, ast.Str):
         return node.s
+    elif isinstance(node, ast.NameConstant):
+        return node.value
     elif isinstance(node, ast.Call):
         funcname = node.func.id
         if funcname == "Fraction":