From a1c44fe804f5251e564a7727f7eb3886034c495e Mon Sep 17 00:00:00 2001
From: Jean-Christophe Jaskula <jjaskula@amazon.com>
Date: Wed, 27 Dec 2023 12:56:18 -0500
Subject: [PATCH] add default_diagram_builder field

---
 src/braket/circuits/circuit.py                            | 3 ++-
 .../braket/circuits/test_ascii_circuit_diagram.py         | 8 ++++----
 .../braket/circuits/test_box_drawing_circuit_diagram.py   | 8 ++++----
 test/unit_tests/braket/circuits/test_circuit.py           | 7 +++++++
 4 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/src/braket/circuits/circuit.py b/src/braket/circuits/circuit.py
index 8075294e5..c3f36b3d3 100644
--- 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
index 8a86427bf..2c27da0d3 100644
--- 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
index 83addfd6c..31cec62f1 100644
--- 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
index 298d7c7d3..fb7dbc7d1 100644
--- 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])