add default_diagram_builder field

amazon-braket · Dec 27, 2023 · a1c44fe · a1c44fe
1 parent d627658
commit a1c44fe
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 9 deletions.
diff --git a/src/braket/circuits/circuit.py b/src/braket/circuits/circuit.py
@@ -82,6 +82,7 @@ class Circuit:
     """
 
     _ALL_QUBITS = "ALL"  # Flag to indicate all qubits in _qubit_observable_mapping
+    default_diagram_builder = BoxDrawingCircuitDiagram
 
     @classmethod
     def register_subroutine(cls, func: SubroutineCallable) -> None:
@@ -1493,7 +1494,7 @@ def __repr__(self) -> str:
             )
 
     def __str__(self):
-        return self.diagram(BoxDrawingCircuitDiagram)
+        return self.diagram(self.default_diagram_builder)
 
     def __eq__(self, other):
         if isinstance(other, Circuit):

diff --git a/test/unit_tests/braket/circuits/test_ascii_circuit_diagram.py b/test/unit_tests/braket/circuits/test_ascii_circuit_diagram.py
@@ -26,6 +26,10 @@
 from braket.pulse import Frame, Port, PulseSequence
 
 
+def _assert_correct_diagram(circ, expected):
+    assert AsciiCircuitDiagram.build_diagram(circ) == "\n".join(expected)
+
+
 def test_empty_circuit():
     assert AsciiCircuitDiagram.build_diagram(Circuit()) == ""
 
@@ -787,10 +791,6 @@ def test_pulse_gate_multi_qubit_circuit():
     _assert_correct_diagram(circ, expected)
 
 
-def _assert_correct_diagram(circ, expected):
-    assert AsciiCircuitDiagram.build_diagram(circ) == "\n".join(expected)
-
-
 def test_circuit_with_nested_target_list():
     circ = (
         Circuit()

diff --git a/test/unit_tests/braket/circuits/test_box_drawing_circuit_diagram.py b/test/unit_tests/braket/circuits/test_box_drawing_circuit_diagram.py
@@ -26,6 +26,10 @@
 from braket.pulse import Frame, Port, PulseSequence
 
 
+def _assert_correct_diagram(circ, expected):
+    assert BoxDrawingCircuitDiagram.build_diagram(circ) == "\n".join(expected)
+
+
 def test_empty_circuit():
     assert BoxDrawingCircuitDiagram.build_diagram(Circuit()) == ""
 
@@ -865,10 +869,6 @@ def test_pulse_gate_multi_qubit_circuit():
     _assert_correct_diagram(circ, expected)
 
 
-def _assert_correct_diagram(circ, expected):
-    assert BoxDrawingCircuitDiagram.build_diagram(circ) == "\n".join(expected)
-
-
 def test_circuit_with_nested_target_list():
     circ = (
         Circuit()

diff --git a/test/unit_tests/braket/circuits/test_circuit.py b/test/unit_tests/braket/circuits/test_circuit.py
@@ -18,6 +18,7 @@
 
 import braket.ir.jaqcd as jaqcd
 from braket.circuits import (
+    AsciiCircuitDiagram,
     BoxDrawingCircuitDiagram,
     Circuit,
     FreeParameter,
@@ -183,6 +184,12 @@ def test_str(h):
     assert str(h) == expected
 
 
+def test_change_diagram_builder(h):
+    Circuit.default_diagram_builder = AsciiCircuitDiagram
+    expected = AsciiCircuitDiagram.build_diagram(h)
+    assert str(h) == expected
+
+
 def test_equality():
     circ_1 = Circuit().h(0).probability([0, 1])
     circ_2 = Circuit().h(0).probability([0, 1])