diff --git a/quantumflow/dagcircuit.py b/quantumflow/dagcircuit.py index b0b7dc9..60d87d7 100644 --- a/quantumflow/dagcircuit.py +++ b/quantumflow/dagcircuit.py @@ -253,19 +253,41 @@ def __iter__(self) -> Iterator[Operation]: if not isinstance(elem, In) and not isinstance(elem, Out): yield elem - # DOCME TESTME def next_element(self, elem: Operation, qubit: Optional[Qubit] = None) -> Operation: + """Return the next element in the DAG following the given element. + + Args: + elem: The operation to find the successor of. + qubit: If specified, follow only the edge for this qubit. + + Returns: + The next operation in the DAG. + + Raises: + KeyError: If no next element is found. + """ for _, node, key in self.graph.edges(elem, keys=True): if qubit is None or key == qubit: return node - assert False # Insanity check # FIXME, raise exception + raise KeyError(f"No next element found for {elem!r} on qubit {qubit}") - # DOCME TESTME def prev_element(self, elem: Operation, qubit: Optional[Qubit] = None) -> Operation: + """Return the previous element in the DAG before the given element. + + Args: + elem: The operation to find the predecessor of. + qubit: If specified, follow only the edge for this qubit. + + Returns: + The previous operation in the DAG. + + Raises: + KeyError: If no previous element is found. + """ for node, _, key in self.graph.in_edges(elem, keys=True): if qubit is None or key == qubit: return node - assert False # Insanity check # FIXME, raise exception + raise KeyError(f"No previous element found for {elem!r} on qubit {qubit}") def next_edges(self, elem: Operation) -> List[Tuple[Any, Any, Qubit]]: qubits = elem.qubits