clean up test conditions

amazon-braket · Nov 26, 2023 · dfdecfe · dfdecfe
1 parent ef91d3d
commit dfdecfe
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 5 deletions.
diff --git a/src/braket/circuits/ascii_circuit_diagram.py b/src/braket/circuits/ascii_circuit_diagram.py
@@ -140,8 +140,11 @@ def _ascii_group_items(
             ):
                 continue
 
-            # TODO: We should use isinstance(item.operator, GPhase)
-            if isinstance(item, Instruction) and item.operator.__class__.__name__ == "GPhase":
+            if (
+                isinstance(item, Instruction)
+                and isinstance(item.operator, Gate)
+                and item.operator.name == "GPhase"
+            ):
                 qubit_range = QubitSet()
             elif (isinstance(item, ResultType) and not item.target) or (
                 isinstance(item, Instruction) and isinstance(item.operator, CompilerDirective)
@@ -162,7 +165,9 @@ def _ascii_group_items(
                 instr_group = group[1]
                 # Take into account overlapping multi-qubit gates
                 if not qubits_added.intersection(set(qubit_range)) or (
-                    isinstance(item, Instruction) and item.operator.__class__.__name__ == "GPhase"
+                    isinstance(item, Instruction)
+                    and isinstance(item.operator, Gate)
+                    and item.operator.name == "GPhase"
                 ):
                     instr_group.append(item)
                     qubits_added.update(qubit_range)
@@ -294,7 +299,11 @@ def _ascii_diagram_column(
                 num_after = len(circuit_qubits) - 1
                 after = ["|"] * (num_after - 1) + ([marker] if num_after else [])
                 ascii_symbols = [ascii_symbol] + after
-            elif isinstance(item, Instruction) and item.operator.__class__.__name__ == "GPhase":
+            elif (
+                isinstance(item, Instruction)
+                and isinstance(item.operator, Gate)
+                and item.operator.name == "GPhase"
+            ):
                 target_qubits = circuit_qubits
                 control_qubits = QubitSet()
                 target_and_control = QubitSet()

diff --git a/src/braket/circuits/moments.py b/src/braket/circuits/moments.py
@@ -19,6 +19,7 @@
 from typing import Any, NamedTuple, Union
 
 from braket.circuits.compiler_directive import CompilerDirective
+from braket.circuits.gate import Gate
 from braket.circuits.instruction import Instruction
 from braket.circuits.noise import Noise
 from braket.registers.qubit import Qubit
@@ -184,7 +185,7 @@ def _add(self, instruction: Instruction, noise_index: int = 0) -> None:
             self._time_all_qubits = time
         elif isinstance(operator, Noise):
             self.add_noise(instruction)
-        elif operator.__class__.__name__ == "GPhase":
+        elif isinstance(operator, Gate) and operator.name == "GPhase":
             time = self._get_qubit_times(self._max_times.keys()) + 1
             self._number_gphase_in_current_moment += 1
             key = MomentsKey(