Add pragma (#35)

* Add pragma

* Return better error message

* Extend and reuse add_statement

* Add test for annotation and ifs

* Finalize if before adding annotations

* Fix finalize if with annotation and no else

* Fix typo

Author: jcjaskula-aws

oqpy/program.py

@@ -55,13 +55,15 @@ class ProgramState:
     """
 
     def __init__(self) -> None:
-        self.body: list[ast.Statement] = []
+        self.body: list[ast.Statement | ast.Pragma] = []
         self.if_clause: Optional[ast.BranchingStatement] = None
         self.annotations: list[ast.Annotation] = []
 
     def add_if_clause(self, condition: ast.Expression, if_clause: list[ast.Statement]) -> None:
+        if_clause_annotations, self.annotations = self.annotations, []
         self.finalize_if_clause()
         self.if_clause = ast.BranchingStatement(condition, if_clause, [])
+        self.if_clause.annotations = if_clause_annotations
 
     def add_else_clause(self, else_clause: list[ast.Statement]) -> None:
         if self.if_clause is None:
@@ -74,12 +76,15 @@ def finalize_if_clause(self) -> None:
             if_clause, self.if_clause = self.if_clause, None
             self.add_statement(if_clause)
 
-    def add_statement(self, stmt: ast.Statement) -> None:
-        assert isinstance(stmt, ast.Statement)
-        self.finalize_if_clause()
-        if self.annotations:
+    def add_statement(self, stmt: ast.Statement | ast.Pragma) -> None:
+        # This function accepts Statement and Pragma even though
+        # it seems to conflict with the definition of ast.Program.
+        # Issue raised in https://github.com/openqasm/openqasm/issues/468
+        assert isinstance(stmt, (ast.Statement, ast.Pragma))
+        if isinstance(stmt, ast.Statement) and self.annotations:
             stmt.annotations = self.annotations + list(stmt.annotations)
             self.annotations = []
+        self.finalize_if_clause()
         self.body.append(stmt)
 
 
@@ -457,6 +462,13 @@ def measure(
         )
         return self
 
+    def pragma(self, command: str) -> Program:
+        """Add a pragma instruction."""
+        if len(self.stack) != 1:
+            raise RuntimeError("Pragmas must be global")
+        self._add_statement(ast.Pragma(command))
+        return self
+
     def _do_assignment(self, var: AstConvertible, op: str, value: AstConvertible) -> None:
         """Helper function for variable assignment operations."""
         if isinstance(var, classical_types.DurationVar):
@@ -537,7 +549,9 @@ def visit_SubroutineDefinition(
         node.body = self.process_statement_list(node.body)
         self.generic_visit(node, context)
 
-    def process_statement_list(self, statements: list[ast.Statement]) -> list[ast.Statement]:
+    def process_statement_list(
+        self, statements: list[ast.Statement | ast.Pragma]
+    ) -> list[ast.Statement | ast.Pragma]:
         new_list = []
         cal_stmts = []
         for stmt in statements:

tests/test_directives.py

@@ -350,21 +350,32 @@ def test_binary_expressions():
     assert prog.to_qasm() == expected
 
 
-def test_measure_reset():
+def test_measure_reset_pragma():
     prog = Program()
     q = PhysicalQubits[0]
     c = BitVar(name="c")
     prog.reset(q)
+    prog.pragma("CLASSIFIER linear")
     prog.measure(q, c)
     prog.measure(q)
+    with oqpy.If(prog, c == 1):
+        with pytest.raises(RuntimeError):
+            prog.pragma("Invalid pragma")
+        prog.gate(q, "x")
+    prog.pragma("LOAD_MEMORY all")
 
     expected = textwrap.dedent(
         """
         OPENQASM 3.0;
         bit c;
         reset $0;
+        pragma CLASSIFIER linear
         c = measure $0;
         measure $0;
+        if (c == 1) {
+            x $0;
+        }
+        pragma LOAD_MEMORY all
         """
     ).strip()
 
@@ -1389,6 +1400,21 @@ def f(prog: Program, x: IntVar) -> IntVar:
     prog.annotate("first-invocation")
     prog.do_expression(f(prog, i))
 
+    prog.annotate("annotation-before-if")
+    with If(prog, i != 0):
+        prog.annotate("annotation-in-if")
+        prog.gate(q1, "x")
+    with oqpy.Else(prog):
+        prog.annotate(("annotation-in-else"))
+        prog.delay(make_duration(1e-8), q1)
+    prog.annotate("annotation-after-if")
+
+    prog.annotate("annotation-no-else-before-if")
+    with If(prog, i != 0):
+        prog.annotate("annotation-no-else-in-if")
+        prog.gate(q1, "x")
+    prog.annotate("annotation-no-else-after-if")
+
     prog.annotate("make-for-loop", "with additional info")
     with ForIn(prog, range(1, 1001), "shot") as shot:
         prog.annotate("declaring_j")
@@ -1433,6 +1459,21 @@ def f(int[32] x) -> int[32] {
         qubit q1;
         @first-invocation
         f(i);
+        @annotation-before-if
+        if (i != 0) {
+        @annotation-in-if
+            x q1;
+        } else {
+        @annotation-in-else
+            delay[10.0ns] q1;
+        }
+        @annotation-after-if
+        @annotation-no-else-before-if
+        if (i != 0) {
+        @annotation-no-else-in-if
+            x q1;
+        }
+        @annotation-no-else-after-if
         @make-for-loop with additional info
         for int shot in [1:1000] {
         @declaring_j