From bb4ad5fdce65a2fa20ea2c6f595916792feac324 Mon Sep 17 00:00:00 2001
From: Renato Mello <renato.msf@gmail.com>
Date: Tue, 24 Sep 2024 14:52:56 +0400
Subject: [PATCH 1/7] remove old `draw` and rename `display` as `draw`

---
 src/qibo/models/circuit.py | 22 ----------------------
 1 file changed, 22 deletions(-)

diff --git a/src/qibo/models/circuit.py b/src/qibo/models/circuit.py
index 32be6d9dc6..ce6adf1af4 100644
--- a/src/qibo/models/circuit.py
+++ b/src/qibo/models/circuit.py
@@ -1367,28 +1367,6 @@ def __str__(self):
     def draw(self, line_wrap: int = 70, legend: bool = False):
         """Draw text circuit using unicode symbols.
 
-        Args:
-            line_wrap (int, optional): maximum number of characters per line. This option
-                split the circuit text diagram in chunks of line_wrap characters.
-                Defaults to :math:`70`.
-            legend (bool, optional): If ``True`` prints a legend below the circuit for
-                callbacks and channels. Defaults to ``False``.
-
-        Returns:
-            String containing text circuit diagram.
-        """
-        qibo.config.log.warning(
-            "Starting on qibo 0.2.13, ``Circuit.draw`` will work in-place. "
-            + "The in-place method is currently implemented as ``Circuit.display``, but "
-            + "will be renamed as ``Circuit.draw`` on release 0.2.13. "
-            + "In release 0.2.12, the in-place display of circuits is accessible as "
-            + "``Circuit.display``."
-        )
-        return self.diagram(line_wrap, legend)
-
-    def display(self, line_wrap: int = 70, legend: bool = False):
-        """Draw text circuit using unicode symbols.
-
         Args:
             line_wrap (int, optional): maximum number of characters per line. This option
                 split the circuit text diagram in chunks of line_wrap characters.

From e3853202bae3b79dffdde8645ba892295e2a1fb0 Mon Sep 17 00:00:00 2001
From: Renato Mello <renato.msf@gmail.com>
Date: Tue, 24 Sep 2024 15:06:01 +0400
Subject: [PATCH 2/7] modifications

---
 doc/source/code-examples/advancedexamples.rst          |  6 +++---
 doc/source/code-examples/examples.rst                  |  2 +-
 examples/adiabatic_qml/adiabatic-qml.ipynb             |  2 +-
 examples/anomaly_detection/test.py                     |  2 +-
 examples/anomaly_detection/train.py                    |  2 +-
 .../qibo-draw-circuit-matplotlib.ipynb                 | 10 +++++-----
 examples/qcnn_classifier/qcnn_demo.ipynb               |  2 +-
 examples/qfiae/qfiae_demo.ipynb                        |  4 ++--
 src/qibo/gates/abstract.py                             |  2 +-
 src/qibo/gates/measurements.py                         |  2 +-
 tests/test_measurements.py                             |  2 +-
 tests/test_models_circuit.py                           |  6 +++---
 tests/test_models_circuit_fuse.py                      |  2 +-
 tests/test_noise.py                                    |  2 +-
 14 files changed, 23 insertions(+), 23 deletions(-)

diff --git a/doc/source/code-examples/advancedexamples.rst b/doc/source/code-examples/advancedexamples.rst
index 50ebdef539..e6740b5141 100644
--- a/doc/source/code-examples/advancedexamples.rst
+++ b/doc/source/code-examples/advancedexamples.rst
@@ -1151,7 +1151,7 @@ Additionally, one can also pass single-qubit readout error probabilities (`reado
     )
 
     print("raw circuit:")
-    print(circuit.draw())
+    circuit.draw()
 
     parameters = {
         "t1": {"0": 250*1e-06, "1": 240*1e-06},
@@ -1168,7 +1168,7 @@ Additionally, one can also pass single-qubit readout error probabilities (`reado
     noisy_circuit = noise_model.apply(circuit)
 
     print("noisy circuit:")
-    print(noisy_circuit.draw())
+    noisy_circuit.draw()
 
 .. testoutput::
    :hide:
@@ -1242,7 +1242,7 @@ Let's see how to use them. For starters, let's define a dummy circuit with some
    circ.add(gates.M(*range(nqubits)))
 
    # visualize the circuit
-   print(circ.draw())
+   circ.draw()
 
    #  q0: ─RZ─RX─RZ─RX─RZ─o────o────────M─
    #  q1: ─RZ─RX─RZ─RX─RZ─X─RZ─X─o────o─M─
diff --git a/doc/source/code-examples/examples.rst b/doc/source/code-examples/examples.rst
index 081e8fc71d..ca5bdc190c 100644
--- a/doc/source/code-examples/examples.rst
+++ b/doc/source/code-examples/examples.rst
@@ -310,7 +310,7 @@ For example
     from qibo.models import QFT
 
     c = QFT(5)
-    print(c.draw())
+    c.draw()
     # Prints
     '''
     q0: ─H─U1─U1─U1─U1───────────────────────────x───
diff --git a/examples/adiabatic_qml/adiabatic-qml.ipynb b/examples/adiabatic_qml/adiabatic-qml.ipynb
index 0ee08d3d32..66d57a4d19 100644
--- a/examples/adiabatic_qml/adiabatic-qml.ipynb
+++ b/examples/adiabatic_qml/adiabatic-qml.ipynb
@@ -568,7 +568,7 @@
     "circ1 = rotcirc.rotations_circuit(t=0.1)\n",
     "circ2 = rotcirc.rotations_circuit(t=0.8)\n",
     "\n",
-    "print(f\"Circuit diagram: {circ1.draw()}\")\n",
+    "print(f\"Circuit diagram: {str(circ1)}\")\n",
     "print(f\"\\nCirc1 params: {circ1.get_parameters()}\")\n",
     "print(f\"\\nCirc2 params: {circ2.get_parameters()}\")"
    ]
diff --git a/examples/anomaly_detection/test.py b/examples/anomaly_detection/test.py
index 043db58a0c..6104fa6294 100644
--- a/examples/anomaly_detection/test.py
+++ b/examples/anomaly_detection/test.py
@@ -99,7 +99,7 @@ def compute_loss_test(encoder, vector):
     encoder_test = make_encoder(n_qubits, n_layers, trained_params, q_compression)
     encoder_test.compile()
     print("Circuit model summary")
-    print(encoder_test.draw())
+    encoder_test.draw()
 
     print("Computing losses...")
     # Compute loss for standard data
diff --git a/examples/anomaly_detection/train.py b/examples/anomaly_detection/train.py
index 0f662a01f3..0406a1f122 100644
--- a/examples/anomaly_detection/train.py
+++ b/examples/anomaly_detection/train.py
@@ -123,7 +123,7 @@ def train_step(batch_size, encoder, params, dataset):
     # Create and print encoder circuit
     encoder = make_encoder(n_qubits, n_layers, params, q_compression)
     print("Circuit model summary")
-    print(encoder.draw())
+    encoder.draw()
 
     # Define optimizer parameters
     steps_for_epoch = math.ceil(train_size / batch_size)
diff --git a/examples/circuit-draw-mpl/qibo-draw-circuit-matplotlib.ipynb b/examples/circuit-draw-mpl/qibo-draw-circuit-matplotlib.ipynb
index d493556c1e..a9ae2410e1 100644
--- a/examples/circuit-draw-mpl/qibo-draw-circuit-matplotlib.ipynb
+++ b/examples/circuit-draw-mpl/qibo-draw-circuit-matplotlib.ipynb
@@ -101,7 +101,7 @@
     "    \n",
     "ansatz.add((gates.RY(q, theta=0) for q in range(nqubits)))\n",
     "ansatz.add(gates.M(qubit) for qubit in range(2))\n",
-    "print(ansatz.draw())"
+    "ansatz.draw()"
    ]
   },
   {
@@ -185,7 +185,7 @@
     "c.add(gates.X(1))\n",
     "c.add(gates.X(0))\n",
     "c.add(gates.M(qubit) for qubit in range(2))\n",
-    "print(c.draw())"
+    "c.draw()"
    ]
   },
   {
@@ -206,7 +206,7 @@
     }
    ],
    "source": [
-    "plot_circuit(c);"
+    "plot_circuit(c)"
    ]
   },
   {
@@ -231,7 +231,7 @@
     "c = QFT(5)\n",
     "c.add(gates.M(qubit) for qubit in range(2))\n",
     "\n",
-    "print(c.draw())"
+    "c.draw()"
    ]
   },
   {
@@ -384,7 +384,7 @@
     }
    ],
    "source": [
-    "print(c.fuse().draw())"
+    "c.fuse().draw()"
    ]
   },
   {
diff --git a/examples/qcnn_classifier/qcnn_demo.ipynb b/examples/qcnn_classifier/qcnn_demo.ipynb
index 0f72938bab..6b372e22e1 100644
--- a/examples/qcnn_classifier/qcnn_demo.ipynb
+++ b/examples/qcnn_classifier/qcnn_demo.ipynb
@@ -240,7 +240,7 @@
    "source": [
     "test = QuantumCNN(nqubits=4, nlayers=1, nclasses=2)\n",
     "testcircuit = test._circuit\n",
-    "print(testcircuit.draw())"
+    "testcircuit.draw()"
    ]
   },
   {
diff --git a/examples/qfiae/qfiae_demo.ipynb b/examples/qfiae/qfiae_demo.ipynb
index 50448ccd75..04f118d5ba 100644
--- a/examples/qfiae/qfiae_demo.ipynb
+++ b/examples/qfiae/qfiae_demo.ipynb
@@ -655,7 +655,7 @@
    ],
    "source": [
     "circuit_a_prob = create_circuit_a(n_qubits = 4,b_max = x_max_int, b_min = x_min_int)\n",
-    "print(circuit_a_prob.draw())"
+    "circuit_a_prob.draw()"
    ]
   },
   {
@@ -753,7 +753,7 @@
     }
    ],
    "source": [
-    "print(create_circuit_q(n_qubits=4,circuit_a=circuit_a_prob).draw())"
+    "create_circuit_q(n_qubits=4,circuit_a=circuit_a_prob).draw()"
    ]
   },
   {
diff --git a/src/qibo/gates/abstract.py b/src/qibo/gates/abstract.py
index 250c309934..3429cb3c8a 100644
--- a/src/qibo/gates/abstract.py
+++ b/src/qibo/gates/abstract.py
@@ -268,7 +268,7 @@ def on_qubits(self, qubit_map) -> "Gate":
                 c.add(gates.CNOT(2, 3).on_qubits({2: 3, 3: 0})) # equivalent to gates.CNOT(3, 0)
                 c.add(gates.CNOT(2, 3).on_qubits({2: 1, 3: 3})) # equivalent to gates.CNOT(1, 3)
                 c.add(gates.CNOT(2, 3).on_qubits({2: 2, 3: 1})) # equivalent to gates.CNOT(2, 1)
-                print(c.draw())
+                c.draw()
             .. testoutput::
 
                 q0: ───X─────
diff --git a/src/qibo/gates/measurements.py b/src/qibo/gates/measurements.py
index 64a7a98e50..5e50ecf1c6 100644
--- a/src/qibo/gates/measurements.py
+++ b/src/qibo/gates/measurements.py
@@ -247,7 +247,7 @@ def on_qubits(self, qubit_map) -> "Gate":
                 c = models.Circuit(3)
                 c.add(measurement.on_qubits({0: 0, 1: 2}))
                 assert c.queue[0].result is measurement.result
-                print(c.draw())
+                c.draw()
             .. testoutput::
 
                 q0: ─M─
diff --git a/tests/test_measurements.py b/tests/test_measurements.py
index 5a8e6a66e6..34d8851dc8 100644
--- a/tests/test_measurements.py
+++ b/tests/test_measurements.py
@@ -443,7 +443,7 @@ def test_measurement_basis_list(backend):
     result = backend.execute_circuit(c, nshots=100)
     assert result.frequencies() == {"0011": 100}
     assert (
-        c.draw()
+        str(c)
         == """q0: ─H─H───M─
 q1: ───────M─
 q2: ─X─H─H─M─
diff --git a/tests/test_models_circuit.py b/tests/test_models_circuit.py
index f05bc6c979..187da44a6a 100644
--- a/tests/test_models_circuit.py
+++ b/tests/test_models_circuit.py
@@ -634,7 +634,7 @@ def test_circuit_draw():
     circuit.add(gates.SWAP(0, 4))
     circuit.add(gates.SWAP(1, 3))
 
-    assert circuit.draw() == ref
+    assert str(circuit) == ref
 
 
 def test_circuit_wire_names_errors():
@@ -667,7 +667,7 @@ def test_circuit_draw_wire_names():
     circuit.add(gates.SWAP(0, 4))
     circuit.add(gates.SWAP(1, 3))
 
-    assert circuit.draw() == ref
+    assert str(circuit) == ref
 
 
 def test_circuit_draw_line_wrap():
@@ -891,7 +891,7 @@ def test_circuit_draw_labels():
             circuit.add(gate)
     circuit.add(gates.SWAP(0, 4))
     circuit.add(gates.SWAP(1, 3))
-    assert circuit.draw() == ref
+    assert str(circuit) == ref
 
 
 def test_circuit_draw_names(capsys):
diff --git a/tests/test_models_circuit_fuse.py b/tests/test_models_circuit_fuse.py
index 6ba1ef729b..22e9c74f0e 100644
--- a/tests/test_models_circuit_fuse.py
+++ b/tests/test_models_circuit_fuse.py
@@ -236,4 +236,4 @@ def test_fused_gate_draw():
     circuit.add(gates.SWAP(0, 4))
     circuit.add(gates.SWAP(1, 3))
     circuit = circuit.fuse()
-    assert circuit.draw() == ref
+    assert str(circuit) == ref
diff --git a/tests/test_noise.py b/tests/test_noise.py
index 325ec95e2d..c7518a3a40 100644
--- a/tests/test_noise.py
+++ b/tests/test_noise.py
@@ -758,7 +758,7 @@ def test_ibmq_noise(
 
     noisy_circuit_target = noise_model_target.apply(circuit)
 
-    assert noisy_circuit.draw() == noisy_circuit_target.draw()
+    assert str(noisy_circuit) == str(noisy_circuit_target)
 
     backend.set_seed(2024)
     state = backend.execute_circuit(noisy_circuit, nshots=10)

From f99b52a586e43a30acc4c005cfa0259f8466207b Mon Sep 17 00:00:00 2001
From: Renato Mello <renato.msf@gmail.com>
Date: Tue, 24 Sep 2024 16:01:18 +0400
Subject: [PATCH 3/7] fix tests

---
 tests/test_models_circuit.py | 224 ++++++++++++++++-------------------
 1 file changed, 101 insertions(+), 123 deletions(-)

diff --git a/tests/test_models_circuit.py b/tests/test_models_circuit.py
index 187da44a6a..70ae69e708 100644
--- a/tests/test_models_circuit.py
+++ b/tests/test_models_circuit.py
@@ -2,6 +2,7 @@
 
 from collections import Counter
 
+import numpy as np
 import pytest
 
 from qibo import Circuit, gates
@@ -500,8 +501,6 @@ def test_circuit_with_pauli_noise(measurements, noise_map):
 @pytest.mark.parametrize("include_not_trainable", [True, False])
 @pytest.mark.parametrize("format", ["list", "dict", "flatlist"])
 def test_get_parameters(trainable, include_not_trainable, format):
-    import numpy as np
-
     matrix = np.random.random((2, 2))
     c = Circuit(3)
     c.add(gates.RX(0, theta=0.123))
@@ -654,10 +653,10 @@ def test_circuit_draw_wire_names():
     """Test circuit text draw."""
     ref = (
         "a    : ─H─U1─U1─U1─U1───────────────────────────x───\n"
-        "b    : ───o──|──|──|──H─U1─U1─U1────────────────|─x─\n"
-        "hello: ──────o──|──|────o──|──|──H─U1─U1────────|─|─\n"
-        "1    : ─────────o──|───────o──|────o──|──H─U1───|─x─\n"
-        "q4   : ────────────o──────────o───────o────o──H─x───"
+        + "b    : ───o──|──|──|──H─U1─U1─U1────────────────|─x─\n"
+        + "hello: ──────o──|──|────o──|──|──H─U1─U1────────|─|─\n"
+        + "1    : ─────────o──|───────o──|────o──|──H─U1───|─x─\n"
+        + "q4   : ────────────o──────────o───────o────o──H─x───"
     )
     circuit = Circuit(5, wire_names=["a", "b", "hello", "1", "q4"])
     for i1 in range(5):
@@ -670,44 +669,42 @@ def test_circuit_draw_wire_names():
     assert str(circuit) == ref
 
 
-def test_circuit_draw_line_wrap():
+def test_circuit_draw_line_wrap(capsys):
     """Test circuit text draw with line wrap."""
     ref_line_wrap_50 = (
         "q0:     ─H─U1─U1─U1─U1───────────────────────────x───I───f ...\n"
-        "q1:     ───o──|──|──|──H─U1─U1─U1────────────────|─x─I───| ...\n"
-        "q2:     ──────o──|──|────o──|──|──H─U1─U1────────|─|─────| ...\n"
-        "q3:     ─────────o──|───────o──|────o──|──H─U1───|─x───M─| ...\n"
-        "q4:     ────────────o──────────o───────o────o──H─x───────f ...\n"
-        "\n"
-        "q0: ... ─o────gf───M─\n"
-        "q1: ... ─U3───|──o─M─\n"
-        "q2: ... ────X─gf─o─M─\n"
-        "q3: ... ────o────o───\n"
-        "q4: ... ────o────X───"
+        + "q1:     ───o──|──|──|──H─U1─U1─U1────────────────|─x─I───| ...\n"
+        + "q2:     ──────o──|──|────o──|──|──H─U1─U1────────|─|─────| ...\n"
+        + "q3:     ─────────o──|───────o──|────o──|──H─U1───|─x───M─| ...\n"
+        + "q4:     ────────────o──────────o───────o────o──H─x───────f ...\n"
+        + "\n"
+        + "q0: ... ─o────gf───M─\n"
+        + "q1: ... ─U3───|──o─M─\n"
+        + "q2: ... ────X─gf─o─M─\n"
+        + "q3: ... ────o────o───\n"
+        + "q4: ... ────o────X───"
     )
 
     ref_line_wrap_30 = (
         "q0:     ─H─U1─U1─U1─U1──────────────── ...\n"
-        "q1:     ───o──|──|──|──H─U1─U1─U1───── ...\n"
-        "q2:     ──────o──|──|────o──|──|──H─U1 ...\n"
-        "q3:     ─────────o──|───────o──|────o─ ...\n"
-        "q4:     ────────────o──────────o────── ...\n"
-        "\n"
-        "q0: ... ───────────x───I───f─o────gf── ...\n"
-        "q1: ... ───────────|─x─I───|─U3───|──o ...\n"
-        "q2: ... ─U1────────|─|─────|────X─gf─o ...\n"
-        "q3: ... ─|──H─U1───|─x───M─|────o────o ...\n"
-        "q4: ... ─o────o──H─x───────f────o────X ...\n"
-        "\n"
-        "q0: ... ─M─\n"
-        "q1: ... ─M─\n"
-        "q2: ... ─M─\n"
-        "q3: ... ───\n"
-        "q4: ... ───"
+        + "q1:     ───o──|──|──|──H─U1─U1─U1───── ...\n"
+        + "q2:     ──────o──|──|────o──|──|──H─U1 ...\n"
+        + "q3:     ─────────o──|───────o──|────o─ ...\n"
+        + "q4:     ────────────o──────────o────── ...\n"
+        + "\n"
+        + "q0: ... ───────────x───I───f─o────gf── ...\n"
+        + "q1: ... ───────────|─x─I───|─U3───|──o ...\n"
+        + "q2: ... ─U1────────|─|─────|────X─gf─o ...\n"
+        + "q3: ... ─|──H─U1───|─x───M─|────o────o ...\n"
+        + "q4: ... ─o────o──H─x───────f────o────X ...\n"
+        + "\n"
+        + "q0: ... ─M─\n"
+        + "q1: ... ─M─\n"
+        + "q2: ... ─M─\n"
+        + "q3: ... ───\n"
+        + "q4: ... ───"
     )
 
-    import numpy as np
-
     circuit = Circuit(5)
     for i1 in range(5):
         circuit.add(gates.H(i1))
@@ -723,58 +720,52 @@ def test_circuit_draw_line_wrap():
     circuit.add(gates.GeneralizedfSim(0, 2, np.eye(2), 0))
     circuit.add(gates.X(4).controlled_by(1, 2, 3))
     circuit.add(gates.M(*range(3)))
-    assert (
-        circuit.draw(
-            line_wrap=50,
-        )
-        == ref_line_wrap_50
-    )
-    assert (
-        circuit.draw(
-            line_wrap=30,
-        )
-        == ref_line_wrap_30
-    )
 
+    circuit.draw(line_wrap=50)
+    out, _ = capsys.readouterr()
+    assert out == ref_line_wrap_50
 
-def test_circuit_draw_line_wrap_names():
+    circuit.draw(line_wrap=30)
+    out, _ = capsys.readouterr()
+    assert out == ref_line_wrap_30
+
+
+def test_circuit_draw_line_wrap_names(capsys):
     """Test circuit text draw with line wrap."""
     ref_line_wrap_50 = (
         "q0:     ─H─U1─U1─U1─U1───────────────────────────x───I───f ...\n"
-        "a :     ───o──|──|──|──H─U1─U1─U1────────────────|─x─I───| ...\n"
-        "q2:     ──────o──|──|────o──|──|──H─U1─U1────────|─|─────| ...\n"
-        "q3:     ─────────o──|───────o──|────o──|──H─U1───|─x───M─| ...\n"
-        "q4:     ────────────o──────────o───────o────o──H─x───────f ...\n"
-        "\n"
-        "q0: ... ─o────gf───M─\n"
-        "a : ... ─U3───|──o─M─\n"
-        "q2: ... ────X─gf─o─M─\n"
-        "q3: ... ────o────o───\n"
-        "q4: ... ────o────X───"
+        + "a :     ───o──|──|──|──H─U1─U1─U1────────────────|─x─I───| ...\n"
+        + "q2:     ──────o──|──|────o──|──|──H─U1─U1────────|─|─────| ...\n"
+        + "q3:     ─────────o──|───────o──|────o──|──H─U1───|─x───M─| ...\n"
+        + "q4:     ────────────o──────────o───────o────o──H─x───────f ...\n"
+        + "\n"
+        + "q0: ... ─o────gf───M─\n"
+        + "a : ... ─U3───|──o─M─\n"
+        + "q2: ... ────X─gf─o─M─\n"
+        + "q3: ... ────o────o───\n"
+        + "q4: ... ────o────X───"
     )
 
     ref_line_wrap_30 = (
         "q0:     ─H─U1─U1─U1─U1──────────────── ...\n"
-        "a :     ───o──|──|──|──H─U1─U1─U1───── ...\n"
-        "q2:     ──────o──|──|────o──|──|──H─U1 ...\n"
-        "q3:     ─────────o──|───────o──|────o─ ...\n"
-        "q4:     ────────────o──────────o────── ...\n"
-        "\n"
-        "q0: ... ───────────x───I───f─o────gf── ...\n"
-        "a : ... ───────────|─x─I───|─U3───|──o ...\n"
-        "q2: ... ─U1────────|─|─────|────X─gf─o ...\n"
-        "q3: ... ─|──H─U1───|─x───M─|────o────o ...\n"
-        "q4: ... ─o────o──H─x───────f────o────X ...\n"
-        "\n"
-        "q0: ... ─M─\n"
-        "a : ... ─M─\n"
-        "q2: ... ─M─\n"
-        "q3: ... ───\n"
-        "q4: ... ───"
+        + "a :     ───o──|──|──|──H─U1─U1─U1───── ...\n"
+        + "q2:     ──────o──|──|────o──|──|──H─U1 ...\n"
+        + "q3:     ─────────o──|───────o──|────o─ ...\n"
+        + "q4:     ────────────o──────────o────── ...\n"
+        + "\n"
+        + "q0: ... ───────────x───I───f─o────gf── ...\n"
+        + "a : ... ───────────|─x─I───|─U3───|──o ...\n"
+        + "q2: ... ─U1────────|─|─────|────X─gf─o ...\n"
+        + "q3: ... ─|──H─U1───|─x───M─|────o────o ...\n"
+        + "q4: ... ─o────o──H─x───────f────o────X ...\n"
+        + "\n"
+        + "q0: ... ─M─\n"
+        + "a : ... ─M─\n"
+        + "q2: ... ─M─\n"
+        + "q3: ... ───\n"
+        + "q4: ... ───"
     )
 
-    import numpy as np
-
     circuit = Circuit(5, wire_names={"q1": "a"})
     for i1 in range(5):
         circuit.add(gates.H(i1))
@@ -790,22 +781,18 @@ def test_circuit_draw_line_wrap_names():
     circuit.add(gates.GeneralizedfSim(0, 2, np.eye(2), 0))
     circuit.add(gates.X(4).controlled_by(1, 2, 3))
     circuit.add(gates.M(*range(3)))
-    assert (
-        circuit.draw(
-            line_wrap=50,
-        )
-        == ref_line_wrap_50
-    )
-    assert (
-        circuit.draw(
-            line_wrap=30,
-        )
-        == ref_line_wrap_30
-    )
+
+    circuit.draw(line_wrap=50)
+    out, _ = capsys.readouterr()
+    assert out == ref_line_wrap_50
+
+    circuit.draw(line_wrap=30)
+    out, _ = capsys.readouterr()
+    assert out == ref_line_wrap_30
 
 
 @pytest.mark.parametrize("legend", [True, False])
-def test_circuit_draw_channels(legend):
+def test_circuit_draw_channels(capsys, legend):
     """Check that channels are drawn correctly."""
 
     circuit = Circuit(2, density_matrix=True)
@@ -823,27 +810,24 @@ def test_circuit_draw_channels(legend):
     circuit.add(gates.CNOT(0, 1))
     circuit.add(gates.DepolarizingChannel((1,), 0.1))
 
-    ref = "q0: ─H─PN─o─PN─o─D─o─D─o───\n" "q1: ─H─PN─X─PN─X─D─X───X─D─"
+    ref = "q0: ─H─PN─o─PN─o─D─o─D─o───\n" + "q1: ─H─PN─X─PN─X─D─X───X─D─"
 
     if legend:
         ref += (
             "\n\n Legend for callbacks and channels: \n"
-            "| Gate                | Symbol   |\n"
-            "|---------------------+----------|\n"
-            "| DepolarizingChannel | D        |\n"
-            "| PauliNoiseChannel   | PN       |"
+            + "| Gate                | Symbol   |\n"
+            + "|---------------------+----------|\n"
+            + "| DepolarizingChannel | D        |\n"
+            + "| PauliNoiseChannel   | PN       |"
         )
 
-    assert (
-        circuit.draw(
-            legend=legend,
-        )
-        == ref
-    )
+    circuit.draw(legend=legend)
+    out, _ = capsys.readouterr()
+    assert out == ref
 
 
 @pytest.mark.parametrize("legend", [True, False])
-def test_circuit_draw_callbacks(legend):
+def test_circuit_draw_callbacks(capsys, legend):
     """Check that callbacks are drawn correcly."""
     from qibo.callbacks import EntanglementEntropy
 
@@ -855,32 +839,29 @@ def test_circuit_draw_callbacks(legend):
     c.add(gates.CNOT(0, 1))
     c.add(gates.CallbackGate(entropy))
 
-    ref = "q0: ─EE─H─EE─o─EE─\n" "q1: ─EE───EE─X─EE─"
+    ref = "q0: ─EE─H─EE─o─EE─\n" + "q1: ─EE───EE─X─EE─"
 
     if legend:
         ref += (
             "\n\n Legend for callbacks and channels: \n"
-            "| Gate                | Symbol   |\n"
-            "|---------------------+----------|\n"
-            "| EntanglementEntropy | EE       |"
+            + "| Gate                | Symbol   |\n"
+            + "|---------------------+----------|\n"
+            + "| EntanglementEntropy | EE       |"
         )
 
-    assert (
-        c.draw(
-            legend=legend,
-        )
-        == ref
-    )
+    c.draw(legend=legend)
+    out, _ = capsys.readouterr()
+    assert out == ref
 
 
 def test_circuit_draw_labels():
     """Test circuit text draw."""
     ref = (
         "q0: ─H─G1─G2─G3─G4───────────────────────────x───\n"
-        "q1: ───o──|──|──|──H─G2─G3─G4────────────────|─x─\n"
-        "q2: ──────o──|──|────o──|──|──H─G3─G4────────|─|─\n"
-        "q3: ─────────o──|───────o──|────o──|──H─G4───|─x─\n"
-        "q4: ────────────o──────────o───────o────o──H─x───"
+        + "q1: ───o──|──|──|──H─G2─G3─G4────────────────|─x─\n"
+        + "q2: ──────o──|──|────o──|──|──H─G3─G4────────|─|─\n"
+        + "q3: ─────────o──|───────o──|────o──|──H─G4───|─x─\n"
+        + "q4: ────────────o──────────o───────o────o──H─x───"
     )
     circuit = Circuit(5)
     for i1 in range(5):
@@ -898,10 +879,10 @@ def test_circuit_draw_names(capsys):
     """Test circuit text draw."""
     ref = (
         "q0: ─H─cx─cx─cx─cx───────────────────────────x───\n"
-        "q1: ───o──|──|──|──H─cx─cx─cx────────────────|─x─\n"
-        "q2: ──────o──|──|────o──|──|──H─cx─cx────────|─|─\n"
-        "q3: ─────────o──|───────o──|────o──|──H─cx───|─x─\n"
-        "q4: ────────────o──────────o───────o────o──H─x───"
+        + "q1: ───o──|──|──|──H─cx─cx─cx────────────────|─x─\n"
+        + "q2: ──────o──|──|────o──|──|──H─cx─cx────────|─|─\n"
+        + "q3: ─────────o──|───────o──|────o──|──H─cx───|─x─\n"
+        + "q4: ────────────o──────────o───────o────o──H─x───"
     )
     circuit = Circuit(5)
     for i1 in range(5):
@@ -912,10 +893,7 @@ def test_circuit_draw_names(capsys):
             circuit.add(gate)
     circuit.add(gates.SWAP(0, 4))
     circuit.add(gates.SWAP(1, 3))
-    assert circuit.draw() == ref
-
-    # Testing circuit text draw when ``output_string == False``
-    circuit.display()
+    circuit.draw()
     out, _ = capsys.readouterr()
     assert out == ref
 

From a222a80700864f46fdef126f400b4767aa55d184 Mon Sep 17 00:00:00 2001
From: Renato Mello <renato.msf@gmail.com>
Date: Sun, 29 Sep 2024 16:18:59 +0400
Subject: [PATCH 4/7] add `flush`

---
 src/qibo/models/circuit.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/qibo/models/circuit.py b/src/qibo/models/circuit.py
index ce6adf1af4..e1fec564d1 100644
--- a/src/qibo/models/circuit.py
+++ b/src/qibo/models/circuit.py
@@ -1378,3 +1378,4 @@ def draw(self, line_wrap: int = 70, legend: bool = False):
             String containing text circuit diagram.
         """
         sys.stdout.write(self.diagram(line_wrap, legend))
+        sys.stdout.flush()

From 19a02a60803e817c9e13d52dcf29c4a431a6ebcc Mon Sep 17 00:00:00 2001
From: Renato Mello <renato.msf@gmail.com>
Date: Sun, 29 Sep 2024 16:51:46 +0400
Subject: [PATCH 5/7] add `\n`

---
 src/qibo/models/circuit.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/src/qibo/models/circuit.py b/src/qibo/models/circuit.py
index e1fec564d1..3ac2c8046d 100644
--- a/src/qibo/models/circuit.py
+++ b/src/qibo/models/circuit.py
@@ -1377,5 +1377,4 @@ def draw(self, line_wrap: int = 70, legend: bool = False):
         Returns:
             String containing text circuit diagram.
         """
-        sys.stdout.write(self.diagram(line_wrap, legend))
-        sys.stdout.flush()
+        sys.stdout.write(self.diagram(line_wrap, legend) + "\n")
\ No newline at end of file

From 50b8bfbfb3b59513e779b571f28a45f7be07d548 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Sun, 29 Sep 2024 12:52:15 +0000
Subject: [PATCH 6/7] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 src/qibo/models/circuit.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/qibo/models/circuit.py b/src/qibo/models/circuit.py
index 3ac2c8046d..8bfb423abf 100644
--- a/src/qibo/models/circuit.py
+++ b/src/qibo/models/circuit.py
@@ -1377,4 +1377,4 @@ def draw(self, line_wrap: int = 70, legend: bool = False):
         Returns:
             String containing text circuit diagram.
         """
-        sys.stdout.write(self.diagram(line_wrap, legend) + "\n")
\ No newline at end of file
+        sys.stdout.write(self.diagram(line_wrap, legend) + "\n")

From b740b241b36d91d3f9309d2f45d1f84a36a26144 Mon Sep 17 00:00:00 2001
From: Renato Mello <renato.msf@gmail.com>
Date: Sun, 29 Sep 2024 16:57:53 +0400
Subject: [PATCH 7/7] fix tests

---
 tests/test_models_circuit.py | 16 ++++++++--------
 1 file changed, 8 insertions(+), 8 deletions(-)

diff --git a/tests/test_models_circuit.py b/tests/test_models_circuit.py
index 70ae69e708..8ffa586740 100644
--- a/tests/test_models_circuit.py
+++ b/tests/test_models_circuit.py
@@ -723,11 +723,11 @@ def test_circuit_draw_line_wrap(capsys):
 
     circuit.draw(line_wrap=50)
     out, _ = capsys.readouterr()
-    assert out == ref_line_wrap_50
+    assert out.rstrip("\n") == ref_line_wrap_50
 
     circuit.draw(line_wrap=30)
     out, _ = capsys.readouterr()
-    assert out == ref_line_wrap_30
+    assert out.rstrip("\n") == ref_line_wrap_30
 
 
 def test_circuit_draw_line_wrap_names(capsys):
@@ -784,11 +784,11 @@ def test_circuit_draw_line_wrap_names(capsys):
 
     circuit.draw(line_wrap=50)
     out, _ = capsys.readouterr()
-    assert out == ref_line_wrap_50
+    assert out.rstrip("\n") == ref_line_wrap_50
 
     circuit.draw(line_wrap=30)
     out, _ = capsys.readouterr()
-    assert out == ref_line_wrap_30
+    assert out.rstrip("\n") == ref_line_wrap_30
 
 
 @pytest.mark.parametrize("legend", [True, False])
@@ -823,7 +823,7 @@ def test_circuit_draw_channels(capsys, legend):
 
     circuit.draw(legend=legend)
     out, _ = capsys.readouterr()
-    assert out == ref
+    assert out.rstrip("\n") == ref
 
 
 @pytest.mark.parametrize("legend", [True, False])
@@ -851,7 +851,7 @@ def test_circuit_draw_callbacks(capsys, legend):
 
     c.draw(legend=legend)
     out, _ = capsys.readouterr()
-    assert out == ref
+    assert out.rstrip("\n") == ref
 
 
 def test_circuit_draw_labels():
@@ -872,7 +872,7 @@ def test_circuit_draw_labels():
             circuit.add(gate)
     circuit.add(gates.SWAP(0, 4))
     circuit.add(gates.SWAP(1, 3))
-    assert str(circuit) == ref
+    assert str(circuit).rstrip("\n") == ref
 
 
 def test_circuit_draw_names(capsys):
@@ -895,7 +895,7 @@ def test_circuit_draw_names(capsys):
     circuit.add(gates.SWAP(1, 3))
     circuit.draw()
     out, _ = capsys.readouterr()
-    assert out == ref
+    assert out.rstrip("\n") == ref
 
 
 def test_circuit_draw_error():