forked from M-Labs/artiq
1
0
Fork 0

transforms/fold_constants: support BoolOp

This commit is contained in:
Sebastien Bourdeauducq 2014-10-14 15:15:02 +08:00
parent b6ac4bdbf3
commit 754a06c623
1 changed files with 31 additions and 3 deletions

View File

@ -29,6 +29,12 @@ _ast_binops = {
} }
_ast_boolops = {
ast.Or: lambda x, y: x or y,
ast.And: lambda x, y: x and y
}
class _ConstantFolder(ast.NodeTransformer): class _ConstantFolder(ast.NodeTransformer):
def visit_UnaryOp(self, node): def visit_UnaryOp(self, node):
self.generic_visit(node) self.generic_visit(node)
@ -62,13 +68,35 @@ class _ConstantFolder(ast.NodeTransformer):
return node return node
return ast.copy_location(result, 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): def visit_Call(self, node):
self.generic_visit(node) self.generic_visit(node)
fn = node.func.id fn = node.func.id
constant_ops = { constant_ops = {
"int": int, "int": int,
"int64": int64, "int64": int64,
"round": round, "round": round,
"round64": round64 "round64": round64
} }
if fn in constant_ops: if fn in constant_ops: