fix subscript if tree

dakk · Jul 10, 2024 · f139cdc · f139cdc
1 parent 2b0278f
commit f139cdc
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 7 deletions.
diff --git 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
@@ -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: