forked from M-Labs/artiq
transforms/fold_constants: support BoolOp
This commit is contained in:
parent
b6ac4bdbf3
commit
754a06c623
|
@ -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:
|
||||||
|
|
Loading…
Reference in New Issue