Add support for arguments & return types to defcals (#20)

* Add arguments and return type to defcals

* Modify how arguments are specified

* Support OQPyExpression in defcal

* Add docstring

* Overload more operators with OQPyExpression

* Fix typo in __rpow__

* Improve test coverage

Author: jcjaskula-aws

oqpy/base.py

@@ -60,12 +60,29 @@ def _to_binary(
         """Helper method to produce a binary expression."""
         return OQPyBinaryExpression(ast.BinaryOperator[op_name], first, second)
 
+    @staticmethod
+    def _to_unary(op_name: str, exp: AstConvertible) -> OQPyUnaryExpression:
+        """Helper method to produce a binary expression."""
+        return OQPyUnaryExpression(ast.UnaryOperator[op_name], exp)
+
+    def __pos__(self) -> OQPyExpression:
+        return self
+
+    def __neg__(self) -> OQPyUnaryExpression:
+        return self._to_unary("-", self)
+
     def __add__(self, other: AstConvertible) -> OQPyBinaryExpression:
         return self._to_binary("+", self, other)
 
     def __radd__(self, other: AstConvertible) -> OQPyBinaryExpression:
         return self._to_binary("+", other, self)
 
+    def __sub__(self, other: AstConvertible) -> OQPyBinaryExpression:
+        return self._to_binary("-", self, other)
+
+    def __rsub__(self, other: AstConvertible) -> OQPyBinaryExpression:
+        return self._to_binary("-", other, self)
+
     def __mod__(self, other: AstConvertible) -> OQPyBinaryExpression:
         return self._to_binary("%", self, other)
 
@@ -78,6 +95,18 @@ def __mul__(self, other: AstConvertible) -> OQPyBinaryExpression:
     def __rmul__(self, other: AstConvertible) -> OQPyBinaryExpression:
         return self._to_binary("*", other, self)
 
+    def __truediv__(self, other: AstConvertible) -> OQPyBinaryExpression:
+        return self._to_binary("/", self, other)
+
+    def __rtruediv__(self, other: AstConvertible) -> OQPyBinaryExpression:
+        return self._to_binary("/", other, self)
+
+    def __pow__(self, other: AstConvertible) -> OQPyBinaryExpression:
+        return self._to_binary("**", self, other)
+
+    def __rpow__(self, other: AstConvertible) -> OQPyBinaryExpression:
+        return self._to_binary("**", other, self)
+
     def __eq__(self, other: AstConvertible) -> OQPyBinaryExpression:  # type: ignore[override]
         return self._to_binary("==", self, other)
 
@@ -132,6 +161,23 @@ def _to_oqpy_expression(self) -> HasToAst:
         ...
 
 
+class OQPyUnaryExpression(OQPyExpression):
+    """An expression consisting of one expression preceded by an operator."""
+
+    def __init__(self, op: ast.UnaryOperator, exp: AstConvertible):
+        super().__init__()
+        self.op = op
+        self.exp = exp
+        if isinstance(exp, OQPyExpression):
+            self.type = exp.type
+        else:
+            raise TypeError("exp is an expression")
+
+    def to_ast(self, program: Program) -> ast.UnaryExpression:
+        """Converts the OQpy expression into an ast node."""
+        return ast.UnaryExpression(self.op, to_ast(program, self.exp))
+
+
 class OQPyBinaryExpression(OQPyExpression):
     """An expression consisting of two subexpressions joined by an operator."""
 

oqpy/classical_types.py

@@ -38,6 +38,7 @@
     from oqpy.program import Program
 
 __all__ = [
+    "pi",
     "BoolVar",
     "IntVar",
     "UintVar",
@@ -53,6 +54,7 @@
     "stretch",
     "bool_",
     "bit_",
+    "bit",
     "bit8",
     "convert_range",
     "int_",
@@ -78,24 +80,24 @@
 # subclasses of ``_ClassicalVar`` instead.
 
 
-def int_(size: int) -> ast.IntType:
+def int_(size: int | None = None) -> ast.IntType:
     """Create a sized signed integer type."""
-    return ast.IntType(ast.IntegerLiteral(size))
+    return ast.IntType(ast.IntegerLiteral(size) if size is not None else None)
 
 
-def uint_(size: int) -> ast.UintType:
+def uint_(size: int | None = None) -> ast.UintType:
     """Create a sized unsigned integer type."""
-    return ast.UintType(ast.IntegerLiteral(size))
+    return ast.UintType(ast.IntegerLiteral(size) if size is not None else None)
 
 
-def float_(size: int) -> ast.FloatType:
+def float_(size: int | None = None) -> ast.FloatType:
     """Create a sized floating-point type."""
-    return ast.FloatType(ast.IntegerLiteral(size))
+    return ast.FloatType(ast.IntegerLiteral(size) if size is not None else None)
 
 
-def angle_(size: int) -> ast.AngleType:
+def angle_(size: int | None = None) -> ast.AngleType:
     """Create a sized angle type."""
-    return ast.AngleType(ast.IntegerLiteral(size))
+    return ast.AngleType(ast.IntegerLiteral(size) if size is not None else None)
 
 
 def complex_(size: int) -> ast.ComplexType:
@@ -107,14 +109,15 @@ def complex_(size: int) -> ast.ComplexType:
     return ast.ComplexType(ast.FloatType(ast.IntegerLiteral(size // 2)))
 
 
-def bit_(size: int) -> ast.BitType:
+def bit_(size: int | None = None) -> ast.BitType:
     """Create a sized bit type."""
-    return ast.BitType(ast.IntegerLiteral(size))
+    return ast.BitType(ast.IntegerLiteral(size) if size is not None else None)
 
 
 duration = ast.DurationType()
 stretch = ast.StretchType()
 bool_ = ast.BoolType()
+bit = ast.BitType()
 bit8 = bit_(8)
 int32 = int_(32)
 int64 = int_(64)
@@ -136,6 +139,22 @@ def convert_range(program: Program, item: Union[slice, range]) -> ast.RangeDefin
     )
 
 
+class Identifier(OQPyExpression):
+    """Base class to specify constant symbols."""
+
+    name: str
+
+    def __init__(self, name: str) -> None:
+        self.type = None
+        self.name = name
+
+    def to_ast(self, program: Program) -> ast.Expression:
+        return ast.Identifier(name=self.name)
+
+
+pi = Identifier(name="pi")
+
+
 class _ClassicalVar(Var, OQPyExpression):
     """Base type for variables with classical type.
 

oqpy/program.py

@@ -82,7 +82,9 @@ class Program:
 
     def __init__(self, version: Optional[str] = "3.0") -> None:
         self.stack: list[ProgramState] = [ProgramState()]
-        self.defcals: dict[tuple[tuple[str, ...], str], ast.CalibrationDefinition] = {}
+        self.defcals: dict[
+            tuple[tuple[str, ...], str, tuple[str, ...]], ast.CalibrationDefinition
+        ] = {}
         self.subroutines: dict[str, ast.SubroutineDefinition] = {}
         self.externs: dict[str, ast.ExternDeclaration] = {}
         self.declared_vars: dict[str, Var] = {}
@@ -196,13 +198,17 @@ def _add_subroutine(self, name: str, stmt: ast.SubroutineDefinition) -> None:
         self.subroutines[name] = stmt
 
     def _add_defcal(
-        self, qubit_names: list[str], name: str, stmt: ast.CalibrationDefinition
+        self,
+        qubit_names: list[str],
+        name: str,
+        arguments: list[str],
+        stmt: ast.CalibrationDefinition,
     ) -> None:
         """Register a defcal defined in this program.
 
         Defcals are added to the top of the program upon conversion to ast.
         """
-        self.defcals[(tuple(qubit_names), name)] = stmt
+        self.defcals[(tuple(qubit_names), name, tuple(arguments))] = stmt
 
     def _make_externs_statements(self, auto_encal: bool = False) -> list[ast.ExternDeclaration]:
         """Return a list of extern statements for inclusion at beginning of program.

oqpy/quantum_types.py

@@ -18,11 +18,13 @@
 from __future__ import annotations
 
 import contextlib
-from typing import TYPE_CHECKING, Iterator, Union
+from typing import TYPE_CHECKING, Iterator, Optional, Union
 
 from openpulse import ast
+from openpulse.printer import dumps
 
-from oqpy.base import Var
+from oqpy.base import AstConvertible, Var, to_ast
+from oqpy.classical_types import _ClassicalVar
 
 if TYPE_CHECKING:
     from oqpy.program import Program
@@ -64,30 +66,57 @@ class QubitArray:
 
 
 @contextlib.contextmanager
-def defcal(program: Program, qubits: Union[Qubit, list[Qubit]], name: str) -> Iterator[None]:
+def defcal(
+    program: Program,
+    qubits: Union[Qubit, list[Qubit]],
+    name: str,
+    arguments: Optional[list[AstConvertible]] = None,
+    return_type: Optional[ast.ClassicalType] = None,
+) -> Union[Iterator[None], Iterator[list[_ClassicalVar]], Iterator[_ClassicalVar]]:
     """Context manager for creating a defcal.
 
     .. code-block:: python
 
-        with defcal(program, q1, "X"):
+        with defcal(program, q1, "X", [AngleVar(name="theta"), oqpy.pi/2], oqpy.bit) as theta:
             program.play(frame, waveform)
     """
-    program._push()
-    yield
-    state = program._pop()
-
     if isinstance(qubits, Qubit):
         qubits = [qubits]
+    assert return_type is None or isinstance(return_type, ast.ClassicalType)
+
+    arguments_ast = []
+    variables = []
+    if arguments is not None:
+        for arg in arguments:
+            if isinstance(arg, _ClassicalVar):
+                arguments_ast.append(
+                    ast.ClassicalArgument(type=arg.type, name=ast.Identifier(name=arg.name))
+                )
+                arg._needs_declaration = False
+                variables.append(arg)
+            else:
+                arguments_ast.append(to_ast(program, arg))
+
+    program._push()
+    if len(variables) > 1:
+        yield variables
+    elif len(variables) == 1:
+        yield variables[0]
+    else:
+        yield
+    state = program._pop()
 
     stmt = ast.CalibrationDefinition(
         ast.Identifier(name),
-        [],  # TODO (#52): support arguments
+        arguments_ast,
         [ast.Identifier(q.name) for q in qubits],
-        None,  # TODO (#52): support return type,
+        return_type,
         state.body,
     )
     program._add_statement(stmt)
-    program._add_defcal([qubit.name for qubit in qubits], name, stmt)
+    program._add_defcal(
+        [qubit.name for qubit in qubits], name, [dumps(a) for a in arguments_ast], stmt
+    )
 
 
 @contextlib.contextmanager