fix regression with cache (#95)

openqasm · Sep 26, 2024 · 07d027b · 07d027b
1 parent 6efb6f5
commit 07d027b
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/oqpy/base.py b/oqpy/base.py
@@ -363,7 +363,7 @@ class CachedExpressionConvertible(Protocol):
 
     _oqpy_cache_key: Hashable
 
-    def _to_cached_oqpy_expression(self) -> HasToAst: ...  # pragma: no cover
+    def _to_cached_oqpy_expression(self) -> AstConvertible: ...  # pragma: no cover
 
 
 class OQPyUnaryExpression(OQPyExpression):
@@ -490,8 +490,10 @@ def to_ast(program: Program, item: AstConvertible) -> ast.Expression:
         if item._oqpy_cache_key is None:
             item._oqpy_cache_key = uuid.uuid1()
         if item._oqpy_cache_key not in program.expr_cache:
-            program.expr_cache[item._oqpy_cache_key] = item._to_cached_oqpy_expression()
-        item = program.expr_cache[item._oqpy_cache_key]
+            program.expr_cache[item._oqpy_cache_key] = to_ast(
+                program, item._to_cached_oqpy_expression()
+            )
+        return program.expr_cache[item._oqpy_cache_key]
     if isinstance(item, (complex, np.complexfloating)):
         if item.imag == 0:
             return to_ast(program, item.real)

diff --git a/tests/test_directives.py b/tests/test_directives.py
@@ -1665,6 +1665,7 @@ def _to_cached_oqpy_expression(self):
     assert dur.count == 2
     # This gets computed just once
     assert frame.count == 1
+    assert all(isinstance(v, ast.QASMNode) for v in prog.expr_cache.values())
 
 
 def test_waveform_extern_arg_passing():