diff --git a/qlasskit/ast2ast/astrewriter.py b/qlasskit/ast2ast/astrewriter.py index be005d3..cd0de70 100644 --- a/qlasskit/ast2ast/astrewriter.py +++ b/qlasskit/ast2ast/astrewriter.py @@ -183,6 +183,9 @@ def visit_Subscript(self, node): # noqa: C901 else: tup = node.value + if isinstance(tup, ast.Constant): + tup = tup.value + if not isinstance(tup, ast.Tuple): raise Exception( "Not a tuple in ast2ast visit subscript with not constant node.slice: " @@ -191,13 +194,8 @@ def visit_Subscript(self, node): # noqa: C901 elts = tup.elts - ifex = ast.IfExp( - test=ast.Compare( - left=node.slice, ops=[ast.Eq()], comparators=[ast.Constant(value=0)] - ), - body=elts[0], - orelse=ast.Constant(value=0), - ) + ifex = elts[0] + for i, x in enumerate(elts[1:]): ifex = ast.IfExp( test=ast.Compare( diff --git a/qlasskit/ast2ast/env.py b/qlasskit/ast2ast/env.py index 2af59df..8bbd3e3 100644 --- a/qlasskit/ast2ast/env.py +++ b/qlasskit/ast2ast/env.py @@ -24,6 +24,9 @@ def set_type(self, name, type_annotation): self.types[name] = type_annotation def set_constant(self, name, value): + if isinstance(value, ast.Constant): + return self.set_constant(name, value.value) + self.constants[name] = value if name not in self.types: