Add capability to annotate statements (#34)

* Add capability to annotate statements

* Address review comments

* style/lint fixes

Author: PhilReinhold

oqpy/base.py

@@ -23,7 +23,7 @@
 
 import sys
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Any, Iterable, Union
+from typing import TYPE_CHECKING, Any, Iterable, Sequence, Union
 
 import numpy as np
 from openpulse import ast
@@ -335,3 +335,15 @@ def optional_ast(program: Program, item: AstConvertible | None) -> ast.Expressio
 def map_to_ast(program: Program, items: Iterable[AstConvertible]) -> list[ast.Expression]:
     """Convert a sequence of items into a sequence of ast nodes."""
     return [to_ast(program, item) for item in items]
+
+
+def make_annotations(vals: Sequence[str | tuple[str, str]]) -> list[ast.Annotation]:
+    """Convert strings/tuples of strings into Annotation ast nodes."""
+    anns: list[ast.Annotation] = []
+    for val in vals:
+        if isinstance(val, str):
+            anns.append(ast.Annotation(val))
+        else:
+            keyword, command = val
+            anns.append(ast.Annotation(keyword, command))
+    return anns

oqpy/classical_types.py

@@ -20,14 +20,24 @@
 import functools
 import random
 import string
-from typing import TYPE_CHECKING, Any, Callable, Iterable, Type, TypeVar, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Iterable,
+    Sequence,
+    Type,
+    TypeVar,
+    Union,
+)
 
 from openpulse import ast
 
 from oqpy.base import (
     AstConvertible,
     OQPyExpression,
     Var,
+    make_annotations,
     map_to_ast,
     optional_ast,
     to_ast,
@@ -170,12 +180,14 @@ def __init__(
         init_expression: AstConvertible | None = None,
         name: str | None = None,
         needs_declaration: bool = True,
+        annotations: Sequence[str | tuple[str, str]] = (),
         **type_kwargs: Any,
     ):
         name = name or "".join([random.choice(string.ascii_letters) for _ in range(10)])
         super().__init__(name, needs_declaration=needs_declaration)
         self.type = self.type_cls(**type_kwargs)
         self.init_expression = init_expression
+        self.annotations = annotations
 
     def to_ast(self, program: Program) -> ast.Identifier:
         """Converts the OQpy variable into an ast node."""
@@ -185,7 +197,9 @@ def to_ast(self, program: Program) -> ast.Identifier:
     def make_declaration_statement(self, program: Program) -> ast.Statement:
         """Make an ast statement that declares the OQpy variable."""
         init_expression_ast = optional_ast(program, self.init_expression)
-        return ast.ClassicalDeclaration(self.type, self.to_ast(program), init_expression_ast)
+        stmt = ast.ClassicalDeclaration(self.type, self.to_ast(program), init_expression_ast)
+        stmt.annotations = make_annotations(self.annotations)
+        return stmt
 
 
 class BoolVar(_ClassicalVar):

oqpy/program.py

@@ -23,6 +23,7 @@
 
 from __future__ import annotations
 
+import warnings
 from copy import deepcopy
 from typing import Any, Iterable, Iterator, Optional
 
@@ -56,6 +57,7 @@ class ProgramState:
     def __init__(self) -> None:
         self.body: list[ast.Statement] = []
         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:
         self.finalize_if_clause()
@@ -73,7 +75,11 @@ def finalize_if_clause(self) -> 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:
+            stmt.annotations = self.annotations + list(stmt.annotations)
+            self.annotations = []
         self.body.append(stmt)
 
 
@@ -149,6 +155,8 @@ def _pop(self) -> ProgramState:
         """Close a context by removing the program state from the top stack, and return it."""
         state = self.stack.pop()
         state.finalize_if_clause()
+        if state.annotations:
+            warnings.warn(f"Annotation(s) {state.annotations} not applied to any statement")
         return state
 
     def _add_var(self, var: Var) -> None:
@@ -259,6 +267,8 @@ def to_ast(
 
         assert len(self.stack) == 1
         self._state.finalize_if_clause()
+        if self._state.annotations:
+            warnings.warn(f"Annotation(s) {self._state.annotations} not applied to any statement")
         statements = []
         if include_externs:
             statements += self._make_externs_statements(encal_declarations)
@@ -459,6 +469,11 @@ def _do_assignment(self, var: AstConvertible, op: str, value: AstConvertible) ->
             )
         )
 
+    def do_expression(self, expression: AstConvertible) -> Program:
+        """Add a statement which evaluates a given expression without assigning the output."""
+        self._add_statement(ast.ExpressionStatement(to_ast(self, expression)))
+        return self
+
     def set(
         self,
         var: classical_types._ClassicalVar | classical_types.OQIndexExpression,
@@ -484,6 +499,11 @@ def mod_equals(self, var: classical_types.IntVar, value: AstConvertible) -> Prog
         self._do_assignment(var, "%=", value)
         return self
 
+    def annotate(self, keyword: str, command: Optional[str] = None) -> Program:
+        """Add an annotation to the next statement."""
+        self._state.annotations.append(ast.Annotation(keyword, command))
+        return self
+
 
 class MergeCalStatementsPass(QASMVisitor[None]):
     """Merge adjacent CalibrationStatement ast nodes."""

oqpy/pulse.py

@@ -17,7 +17,7 @@
 
 from __future__ import annotations
 
-from typing import Any
+from typing import Any, Sequence
 
 from openpulse import ast
 
@@ -67,6 +67,7 @@ def __init__(
         phase: AstConvertible = 0,
         name: str | None = None,
         needs_declaration: bool = True,
+        annotations: Sequence[str | tuple[str, str]] = (),
     ):
         if (port is None) != (frequency is None):
             raise ValueError("Must declare both port and frequency or neither.")
@@ -75,4 +76,6 @@ def __init__(
         else:
             assert frequency is not None
             init_expression = OQFunctionCall("newframe", [port, frequency, phase], ast.FrameType)
-        super().__init__(init_expression, name, needs_declaration=needs_declaration)
+        super().__init__(
+            init_expression, name, needs_declaration=needs_declaration, annotations=annotations
+        )

oqpy/quantum_types.py

@@ -18,12 +18,12 @@
 from __future__ import annotations
 
 import contextlib
-from typing import TYPE_CHECKING, Iterator, Optional, Union
+from typing import TYPE_CHECKING, Iterator, Optional, Sequence, Union
 
 from openpulse import ast
 from openpulse.printer import dumps
 
-from oqpy.base import AstConvertible, Var, to_ast
+from oqpy.base import AstConvertible, Var, make_annotations, to_ast
 from oqpy.classical_types import _ClassicalVar
 
 if TYPE_CHECKING:
@@ -36,9 +36,15 @@
 class Qubit(Var):
     """OQpy variable representing a single qubit."""
 
-    def __init__(self, name: str, needs_declaration: bool = True):
+    def __init__(
+        self,
+        name: str,
+        needs_declaration: bool = True,
+        annotations: Sequence[str | tuple[str, str]] = (),
+    ):
         super().__init__(name, needs_declaration=needs_declaration)
         self.name = name
+        self.annotations = annotations
 
     def to_ast(self, prog: Program) -> ast.Expression:
         """Converts the OQpy variable into an ast node."""
@@ -47,7 +53,9 @@ def to_ast(self, prog: Program) -> ast.Expression:
 
     def make_declaration_statement(self, program: Program) -> ast.Statement:
         """Make an ast statement that declares the OQpy variable."""
-        return ast.QubitDeclaration(ast.Identifier(self.name), size=None)
+        decl = ast.QubitDeclaration(ast.Identifier(self.name), size=None)
+        decl.annotations = make_annotations(self.annotations)
+        return decl
 
 
 class PhysicalQubits:

oqpy/subroutines.py

@@ -19,18 +19,18 @@
 
 import functools
 import inspect
-from typing import Callable, get_type_hints
+from typing import Any, Callable, Sequence, TypeVar, get_type_hints
 
 from mypy_extensions import VarArg
 from openpulse import ast
 
 import oqpy.program
-from oqpy.base import AstConvertible, OQPyExpression, to_ast
+from oqpy.base import AstConvertible, OQPyExpression, make_annotations, to_ast
 from oqpy.classical_types import OQFunctionCall, _ClassicalVar
 from oqpy.quantum_types import Qubit
 from oqpy.timing import make_duration
 
-__all__ = ["subroutine", "declare_extern", "declare_waveform_generator"]
+__all__ = ["subroutine", "annotate_subroutine", "declare_extern", "declare_waveform_generator"]
 
 SubroutineParams = [oqpy.Program, VarArg(AstConvertible)]
 
@@ -68,7 +68,11 @@ def increment_variable(int[32] i) {
     """
 
     @functools.wraps(func)
-    def wrapper(program: oqpy.Program, *args: AstConvertible) -> OQFunctionCall:
+    def wrapper(
+        program: oqpy.Program,
+        *args: AstConvertible,
+        annotations: Sequence[str | tuple[str, str]] = (),
+    ) -> OQFunctionCall:
         name = func.__name__
         identifier = ast.Identifier(func.__name__)
         argnames = list(inspect.signature(func).parameters.keys())
@@ -117,13 +121,42 @@ def wrapper(program: oqpy.Program, *args: AstConvertible) -> OQFunctionCall:
             return_type=return_type,
             body=body,
         )
+        stmt.annotations = make_annotations(annotations)
         return OQFunctionCall(identifier, args, return_type, subroutine_decl=stmt)
 
     return wrapper
 
 
+FnType = TypeVar("FnType", bound=Callable[..., Any])
+
+
+def annotate_subroutine(keyword: str, command: str | None = None) -> Callable[[FnType], FnType]:
+    """Add annotation to a subroutine."""
+
+    def annotate_subroutine_decorator(func: FnType) -> FnType:
+        @functools.wraps(func)
+        def wrapper(
+            program: oqpy.Program,
+            *args: AstConvertible,
+            annotations: Sequence[str | tuple[str, str]] = (),
+        ) -> OQFunctionCall:
+            new_ann: str | tuple[str, str]
+            if command is not None:
+                new_ann = keyword, command
+            else:
+                new_ann = keyword
+            return func(program, *args, annotations=list(annotations) + [new_ann])
+
+        return wrapper  # type: ignore[return-value]
+
+    return annotate_subroutine_decorator
+
+
 def declare_extern(
-    name: str, args: list[tuple[str, ast.ClassicalType]], return_type: ast.ClassicalType
+    name: str,
+    args: list[tuple[str, ast.ClassicalType]],
+    return_type: ast.ClassicalType,
+    annotations: Sequence[str | tuple[str, str]] = (),
 ) -> Callable[..., OQFunctionCall]:
     """Declare an extern and return a callable which adds the extern.
 
@@ -145,6 +178,7 @@ def declare_extern(
         [ast.ExternArgument(type=t) for t in arg_types],
         ast.ExternArgument(type=return_type),
     )
+    extern_decl.annotations = make_annotations(annotations)
 
     def call_extern(*call_args: AstConvertible, **call_kwargs: AstConvertible) -> OQFunctionCall:
         new_args = list(call_args) + [None] * len(call_kwargs)
@@ -180,8 +214,10 @@ def call_extern(*call_args: AstConvertible, **call_kwargs: AstConvertible) -> OQ
 
 
 def declare_waveform_generator(
-    name: str, argtypes: list[tuple[str, ast.ClassicalType]]
+    name: str,
+    argtypes: list[tuple[str, ast.ClassicalType]],
+    annotations: Sequence[str | tuple[str, str]] = (),
 ) -> Callable[..., OQFunctionCall]:
     """Create a function which generates waveforms using a specified name and argument signature."""
-    func = declare_extern(name, argtypes, ast.WaveformType())
+    func = declare_extern(name, argtypes, ast.WaveformType(), annotations=annotations)
     return func