Fix assignment to AnnotatedOperation

crates/circuit/src/circuit_data.rs

@@ -15,9 +15,9 @@ use std::cell::RefCell;
 
 use crate::bit_data::BitData;
 use crate::circuit_instruction::{CircuitInstruction, OperationFromPython};
-use crate::imports::{QUANTUM_CIRCUIT, QUBIT};
+use crate::imports::{ANNOTATED_OPERATION, QUANTUM_CIRCUIT, QUBIT};
 use crate::interner::{IndexedInterner, Interner, InternerKey};
-use crate::operations::{Operation, Param, StandardGate};
+use crate::operations::{Operation, OperationRef, Param, StandardGate};
 use crate::packed_instruction::PackedInstruction;
 use crate::parameter_table::{ParameterTable, ParameterTableError, ParameterUse, ParameterUuid};
 use crate::slice::{PySequenceIndex, SequenceIndex};
@@ -1193,10 +1193,26 @@ impl CircuitData {
         for (instruction, bindings) in user_operations {
             // We only put non-standard gates in `user_operations`, so we're not risking creating a
             // previously non-existent Python object.
-            let definition_cache = self.data[instruction]
-                .unpack_py_op(py)?
-                .bind(py)
-                .getattr(_definition_attr)?;
+            let instruction = &self.data[instruction];
+            let definition_cache = if matches!(instruction.op.view(), OperationRef::Operation(_)) {
+                // `Operation` instances don't have a `definition` as part of their interfaces, but
+                // they might be an `AnnotatedOperation`, which is one of our special built-ins.
+                // This should be handled more completely in the user-customisation interface by a
+                // delegating method, but that's not the data model we currently have.
+                let py_op = instruction.unpack_py_op(py)?;
+                let py_op = py_op.bind(py);
+                if !py_op.is_instance(ANNOTATED_OPERATION.get_bound(py))? {
+                    continue;
+                }
+                py_op
+                    .getattr(intern!(py, "base_op"))?
+                    .getattr(_definition_attr)?
+            } else {
+                instruction
+                    .unpack_py_op(py)?
+                    .bind(py)
+                    .getattr(_definition_attr)?
+            };
             if !definition_cache.is_none() {
                 definition_cache.call_method(
                     assign_parameters_attr,

crates/circuit/src/imports.rs

@@ -73,6 +73,8 @@ pub static SINGLETON_CONTROLLED_GATE: ImportOnceCell =
     ImportOnceCell::new("qiskit.circuit.singleton", "SingletonControlledGate");
 pub static CONTROLLED_GATE: ImportOnceCell =
     ImportOnceCell::new("qiskit.circuit", "ControlledGate");
+pub static ANNOTATED_OPERATION: ImportOnceCell =
+    ImportOnceCell::new("qiskit.circuit", "AnnotatedOperation");
 pub static DEEPCOPY: ImportOnceCell = ImportOnceCell::new("copy", "deepcopy");
 pub static QI_OPERATOR: ImportOnceCell = ImportOnceCell::new("qiskit.quantum_info", "Operator");
 pub static WARNINGS_WARN: ImportOnceCell = ImportOnceCell::new("warnings", "warn");

test/python/circuit/test_parameters.py

@@ -627,6 +627,45 @@ def test_expression_partial_binding_zero(self):
         self.assertEqual(fbqc.parameters, set())
         self.assertEqual(float(fbqc.data[0].operation.params[0]), 0)
 
+    def test_assignment_to_annotated_operation(self):
+        """Test that assignments to an ``AnnotatedOperation`` are propagated all the way down."""
+
+        class MyGate(Gate):
+            def __init__(self, param):
+                super().__init__("my_gate", 1, [param])
+                # Eagerly create our definition.
+                _ = self.definition
+
+            def _define(self):
+                self._definition = QuantumCircuit(1, name="my_gate_inner")
+                self._definition.ry(self.params[0], 0)
+
+        theta = Parameter("theta")
+
+        # Sanity check for the test; it won't catch errors if this fails.
+        self.assertEqual(MyGate(theta), MyGate(theta))
+        self.assertNotEqual(MyGate(theta), MyGate(1.23))
+
+        parametric_gate = MyGate(theta)
+        qc = QuantumCircuit(2)
+        qc.append(parametric_gate.control(1, annotated=True), [0, 1], copy=True)
+        assigned = qc.assign_parameters([1.23])
+
+        expected = QuantumCircuit(2)
+        expected.append(MyGate(1.23).control(1, annotated=True), [0, 1])
+        self.assertEqual(assigned, expected)
+        self.assertEqual(
+            assigned.data[0].operation.base_op.definition,
+            expected.data[0].operation.base_op.definition,
+        )
+
+        qc.assign_parameters([1.23], inplace=True)
+        self.assertEqual(qc, expected)
+
+        # Test that the underlying gate was not modified.
+        self.assertEqual(parametric_gate.params, [theta])
+        self.assertEqual(set(parametric_gate.definition.parameters), {theta})
+
     def test_raise_if_assigning_params_not_in_circuit(self):
         """Verify binding parameters which are not present in the circuit raises an error."""
         x = Parameter("x")