Repository: PennyLaneAI/pennylane (python)

Issue:
Preventing printing overflow for long circuits in terminal with `qml.draw_circuit()`
### Feature details

Calling `qml.drawcircuit` and rendering the output to the terminal results in illegible circuits if the circuit width (*i.e.* the number of text columns used) is greater than the column width of the terminal window. This issue is due to the overflow of text used to visualize the circuit onto the next line. For example:

**Current Behavior:**
```
|------------------terminal width ------------------|
 0: ─H─RZ(1.0)──╭C──────────╭X─RX(3.14)─╭C──────────
 ╭X─RX(3.14)──────┤ ⟨Z⟩
 1: ─H─RZ(0.1)──╰X─╭C───────│──RX(3.14)─╰X─╭C───────
 │──RX(3.14)──────┤
 2: ─H─RZ(0.2)─────╰X─╭C────│──RX(3.14)────╰X─╭C────
 │──RX(3.14)──────┤
 3: ─H─RZ(0.3)────────╰X─╭C─│──RX(3.14)───────╰X─╭C─
 │──RX(3.14)──────┤
 4: ─H─RZ(0.4)───────────╰X─╰C─RX(3.14)──────────╰X─
 ╰C─RX(3.14)──────┤
```
A possible solution would be to detect the width of the terminal prior to printing the circuit. With the width detected, the printing of lines could occur with the circuit length modulo the terminal width to facilitate line skipping as needed. This skipping would be repeated to allow the full circuit to be cleanly rendered in the terminal without the overflow issues. For example:

**Alternative Behavior:**
```
|------------------terminal width ------------------|
 0: ─H─RZ(1.0)──╭C──────────╭X─RX(3.14)─╭C────────── 
 1: ─H─RZ(0.1)──╰X─╭C───────│──RX(3.14)─╰X─╭C───────
 2: ─H─RZ(0.2)─────╰X─╭C────│──RX(3.14)────╰X─╭C────
 3: ─H─RZ(0.3)────────╰X─╭C─│──RX(3.14)───────╰X─╭C─
 4: ─H─RZ(0.4)───────────╰X─╰C─RX(3.14)──────────╰X─

 0: ─╭X─RX(3.14)──────┤⟨Z⟩
 1: ─│──RX(3.14)──────┤
 2: ─│──RX(3.14)──────┤
 3: ─│──RX(3.14)──────┤
 4: ─╰C─RX(3.14)──────┤
```

Provided there's interest, I'd be happy to send a pull request with this feature (although I may need assistance testing it on Linux and Windows operating systems). 

### Implementation

To detect the terminal size:

```python
import shutil
cols, _ = shutil.get_terminal_size()
```

Note that `shutil` is part of the standard library and would not require an additional dependency. 

It seems to me this could be integrated fairly trivially into `qml.circuit_drawer.draw()` (below), but perhaps this could be confirmed by a developer more familiar with the code base. 

```python
    def draw(self):
        """Draw the circuit diagram.
        Returns:
            str: The circuit diagram
        """
        rendered_string = ""

        # extract the wire labels as strings and get their maximum length
        wire_names = []
        padding = 0
        for i in range(self.full_representation_grid.num_wires):
            wire_name = str(self.active_wires.labels[i])
            padding = max(padding, len(wire_name))
            wire_names.append(wire_name)

        for i in range(self.full_representation_grid.num_wires):
            # format wire name nicely
            wire = self.full_representation_grid.wire(i)
            s = " {:>" + str(padding) + "}: {}"

            rendered_string += s.format(wire_names[i], 2 * self.charset.WIRE)

            for s in wire:
                rendered_string += s

            rendered_string += "\n"

        for symbol, cache in [
            ("U", self.representation_resolver.unitary_matrix_cache),
            ("H", self.representation_resolver.hermitian_matrix_cache),
            ("M", self.representation_resolver.matrix_cache),
        ]:
            for idx, matrix in enumerate(cache):
                rendered_string += "{}{} =\n{}\n".format(symbol, idx, matrix)

        return rendered_string
```

### How important would you say this feature is?

1: Not important. Would be nice to have.

### Additional information

_No response_

Pull request description:
Temporary fix for Issue #1841
### Before submitting

Please complete the following checklist when submitting a PR:

- [x] All new features must include a unit test.
      If you've fixed a bug or added code that should be tested, add a test to the
      test directory!

- [x] All new functions and code must be clearly commented and documented.
      If you do make documentation changes, make sure that the docs build and
      render correctly by running `make docs`.

- [x] Ensure that the test suite passes, by running `make test`.

- [x] Add a new entry to the `doc/releases/changelog-dev.md` file, summarizing the
      change, and including a link back to the PR.

- [x] The PennyLane source code conforms to
      [PEP8 standards](https://www.python.org/dev/peps/pep-0008/).
      We check all of our code against [Pylint](https://www.pylint.org/).
      To lint modified files, simply `pip install pylint`, and then
      run `pylint pennylane/path/to/file.py`.

When all the above are checked, delete everything above the dashed
line and fill in the pull request template.

------------------------------------------------------------------------------------------------------------

**Context:**

In reference to Issue #1841. From the discussion, this is a temporary fix to the text overflow issue when printing a circuit to the terminal. 

**Description of the Change:**

The code works by allowing the user to set a `max_length` keyword argument when calling `CircuitDrawer.draw()`. Updates have been made to the `draw` method of the `CircuitDrawer` class and functions referencing it to allow for the `max_length` keyword arg. A unit test for the parameter has also been created.

**Benefits:**
Legible circuits in the CLI as we aspire to polynomial length circuits! Step aside, AC0.  

**Possible Drawbacks:**
- This code isn't a long term solution. Its a quick usability fix. From the discussion in Issue #1841, the code may soon be replaced by a refactored `CircuitDrawer` class. 
- I've unfortunately been unable to run the test suite to verify this code and the unit test due to a number of other errors when running the suite. I have developed this with attention to detail around not breaking existing tests. 

**Related GitHub Issues:**
Closes #1841.

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 9947dd30e..2bc308c75 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -35,6 +35,8 @@
 
 * AngleEmbedding now supports `batch_params` decorator. [(#1812)](https://github.com/PennyLaneAI/pennylane/pull/1812)
 
+* CircuitDrawer now supports a `max_length` argument to help prevent text overflows when printing circuits to the CLI. [#1841](https://github.com/PennyLaneAI/pennylane/pull/1841)
+
 <h3>Breaking changes</h3>
 
 * `QuantumTape.trainable_params` now is a list instead of a set. This
@@ -67,6 +69,8 @@
 
 <h3>Bug fixes</h3>
 
+* `qml.CSWAP` and `qml.CRot` now define `control_wires`, and `qml.SWAP`
+
 * `QuantumTape.trainable_params` now is a list instead of a set, making
   it more stable in very rare edge cases.
   [(#1904)](https://github.com/PennyLaneAI/pennylane/pull/1904)
@@ -91,8 +95,7 @@
 
 <h3>Contributors</h3>
 
-This release contains contributions from (in alphabetical order): 
+This release contains contributions from (in alphabetical order):
 
-Guillermo Alonso-Linaje, Olivia Di Matteo, Jalani Kanem, Shumpei Kobayashi, Christina Lee, Alejandro Montanez,
+Guillermo Alonso-Linaje, Benjamin Cordier, Olivia Di Matteo, Jalani Kanem, Shumpei Kobayashi, Christina Lee, Alejandro Montanez,
 Romain Moyard, Maria Schuld, Jay Soni, David Wierichs
-
diff --git a/pennylane/circuit_drawer/circuit_drawer.py b/pennylane/circuit_drawer/circuit_drawer.py
index 947beb6be..6a9a975bb 100644
--- a/pennylane/circuit_drawer/circuit_drawer.py
+++ b/pennylane/circuit_drawer/circuit_drawer.py
@@ -48,6 +48,7 @@ class CircuitDrawer:
         wires (Wires): all wires on the device for which the circuit is drawn
         charset (str, pennylane.circuit_drawer.CharSet, optional): The CharSet that shall be used for drawing.
         show_all_wires (bool): If True, all wires, including empty wires, are printed.
+        max_length (int, optional): Maximum string width (columns) when printing the circuit to the CLI.
     """
 
     def __init__(
@@ -57,6 +58,7 @@ class CircuitDrawer:
         wires,
         charset=None,
         show_all_wires=False,
+        max_length=None,
     ):
         self.operation_grid = Grid(raw_operation_grid)
         self.observable_grid = Grid(raw_observable_grid)
@@ -81,6 +83,9 @@ class CircuitDrawer:
             # as active wires
             self.active_wires = wires.all_wires([wires, self.active_wires])
 
+        # We add a -2 character offset to account for some downstream formatting
+        self.max_length = max_length - 2 if max_length is not None else None
+
         self.representation_resolver = RepresentationResolver(self.charset)
         self.operation_representation_grid = Grid()
         self.observable_representation_grid = Grid()
@@ -357,7 +362,6 @@ class CircuitDrawer:
             s = " {:>" + str(padding) + "}: {}"
 
             rendered_string += s.format(wire_names[i], 2 * self.charset.WIRE)
-
             for s in wire:
                 rendered_string += s
 
@@ -371,4 +375,26 @@ class CircuitDrawer:
             for idx, matrix in enumerate(cache):
                 rendered_string += "{}{} =\n{}\n".format(symbol, idx, matrix)
 
+        # Restrict CLI print width to max_length
+        if self.max_length is not None:
+            wires = rendered_string.split("\n")
+            n_wraps = (len(wires[0]) // self.max_length) + 1
+            rendered_substrings = []
+            for i in range(n_wraps):
+                for wire in wires:
+                    # Print body of circuit. We compute some index and whitespace offsets to ensure the rendered wire matches existing unit tests.
+                    if (self.max_length * (i + 1)) < len(wire):
+                        offset_white_space = (
+                            int((i != 0) or (self.max_length * (i + 1) >= len(wire))) * " "
+                        )
+                        offset_index = int(i != 0)
+                        rendered_substrings.append(
+                            f"{offset_white_space}"
+                            + f"{wire[(i * self.max_length) + offset_index: ((i + 1) * self.max_length) + 1]}"
+                        )
+                    # Last wrap, print the tail of the circuit. We trim the last character of whitespace at the end.
+                    else:
+                        rendered_substrings.append(f" {wire[i * self.max_length + 1:]}"[:-1])
+            rendered_string = "\n".join(rendered_substrings)
+
         return rendered_string
diff --git a/pennylane/circuit_graph.py b/pennylane/circuit_graph.py
index a44c819b4..007e8c488 100644
--- a/pennylane/circuit_graph.py
+++ b/pennylane/circuit_graph.py
@@ -590,13 +590,14 @@ class CircuitGraph:
         self._operations = self.operations_in_order
         self._observables = self.observables_in_order
 
-    def draw(self, charset="unicode", wire_order=None, show_all_wires=False):
+    def draw(self, charset="unicode", wire_order=None, show_all_wires=False, max_length=None):
         """Draw the CircuitGraph as a circuit diagram.
 
         Args:
             charset (str, optional): The charset that should be used. Currently, "unicode" and "ascii" are supported.
             wire_order (Wires or None): the order (from top to bottom) to print the wires of the circuit
             show_all_wires (bool): If True, all wires, including empty wires, are printed.
+            max_length (int, optional): Maximum string width (columns) when printing the circuit to the CLI.
 
         Raises:
             ValueError: If the given charset is not supported
@@ -615,6 +616,7 @@ class CircuitGraph:
             wires=wire_order or self.wires,
             charset=charset,
             show_all_wires=show_all_wires,
+            max_length=max_length,
         )
 
         return drawer.draw()
diff --git a/pennylane/qnode.py b/pennylane/qnode.py
index fa3a5362f..c48037e68 100644
--- a/pennylane/qnode.py
+++ b/pennylane/qnode.py
@@ -765,7 +765,7 @@ class QNode:
         return qml.metric_tensor(self, allow_nonunitary, approx, diag_approx)(*args, **kwargs)
 
     def draw(
-        self, charset="unicode", wire_order=None, show_all_wires=False
+        self, charset="unicode", wire_order=None, show_all_wires=False, max_length=None
     ):  # pylint: disable=unused-argument
         """Draw the quantum tape as a circuit diagram.
 
@@ -775,6 +775,7 @@ class QNode:
             wire_order (Sequence[Any]): The order (from top to bottom) to print the wires of the circuit.
                 If not provided, this defaults to the wire order of the device.
             show_all_wires (bool): If True, all wires, including empty wires, are printed.
+            max_length (int, optional): Maximum string width (columns) when printing the circuit to the CLI.
 
         Raises:
             ValueError: if the given charset is not supported
@@ -860,7 +861,10 @@ class QNode:
             )
 
         return self.qtape.draw(
-            charset=charset, wire_order=wire_order, show_all_wires=show_all_wires
+            charset=charset,
+            wire_order=wire_order,
+            show_all_wires=show_all_wires,
+            max_length=max_length,
         )
 
     @property
diff --git a/pennylane/tape/tape.py b/pennylane/tape/tape.py
index 4b5212adc..81ae202df 100644
--- a/pennylane/tape/tape.py
+++ b/pennylane/tape/tape.py
@@ -1087,7 +1087,7 @@ class QuantumTape(AnnotatedQueue):
 
         return self._specs
 
-    def draw(self, charset="unicode", wire_order=None, show_all_wires=False):
+    def draw(self, charset="unicode", wire_order=None, show_all_wires=False, max_length=None):
         """Draw the quantum tape as a circuit diagram.
 
         Consider the following circuit as an example:
@@ -1115,6 +1115,7 @@ class QuantumTape(AnnotatedQueue):
                 "ascii" are supported.
             wire_order (Sequence[Any]): the order (from top to bottom) to print the wires of the circuit
             show_all_wires (bool): If True, all wires, including empty wires, are printed.
+            max_length (int, optional): Maximum string width (columns) when printing the circuit to the CLI.
 
         Raises:
             ValueError: if the given charset is not supported
@@ -1126,6 +1127,7 @@ class QuantumTape(AnnotatedQueue):
             charset=charset,
             wire_order=wire_order,
             show_all_wires=show_all_wires,
+            max_length=max_length,
         )
 
     def to_openqasm(self, wires=None, rotations=True, measure_all=True):
diff --git a/pennylane/transforms/draw.py b/pennylane/transforms/draw.py
index 5eaf7d6be..d0fa401ce 100644
--- a/pennylane/transforms/draw.py
+++ b/pennylane/transforms/draw.py
@@ -1,3 +1,5 @@
+# pylint: disable=too-many-arguments
+
 # Copyright 2018-2021 Xanadu Quantum Technologies Inc.
 
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -19,7 +21,14 @@ from functools import wraps
 import pennylane as qml
 
 
-def draw(qnode, charset="unicode", wire_order=None, show_all_wires=False, expansion_strategy=None):
+def draw(
+    qnode,
+    charset="unicode",
+    wire_order=None,
+    show_all_wires=False,
+    max_length=None,
+    expansion_strategy=None,
+):
     """Create a function that draws the given qnode.
 
     Args:
@@ -28,6 +37,7 @@ def draw(qnode, charset="unicode", wire_order=None, show_all_wires=False, expans
             "ascii" are supported.
         wire_order (Sequence[Any]): the order (from top to bottom) to print the wires of the circuit
         show_all_wires (bool): If True, all wires, including empty wires, are printed.
+        max_length (int, optional): Maximum string width (columns) when printing the circuit to the CLI.
         expansion_strategy (str): The strategy to use when circuit expansions or decompositions
             are required.
 
@@ -121,13 +131,21 @@ def draw(qnode, charset="unicode", wire_order=None, show_all_wires=False, expans
 
         if tapes is not None:
             res = [
-                t.draw(charset=charset, wire_order=_wire_order, show_all_wires=show_all_wires)
+                t.draw(
+                    charset=charset,
+                    wire_order=_wire_order,
+                    show_all_wires=show_all_wires,
+                    max_length=max_length,
+                )
                 for t in tapes[0]
             ]
             return "\n".join(res)
 
         return qnode.qtape.draw(
-            charset=charset, wire_order=_wire_order, show_all_wires=show_all_wires
+            charset=charset,
+            wire_order=_wire_order,
+            show_all_wires=show_all_wires,
+            max_length=max_length,
         )
 
     return wrapper


Test patch:
diff --git a/tests/circuit_drawer/test_circuit_drawer.py b/tests/circuit_drawer/test_circuit_drawer.py
index 744b278e4..3d6b24c9c 100644
--- a/tests/circuit_drawer/test_circuit_drawer.py
+++ b/tests/circuit_drawer/test_circuit_drawer.py
@@ -767,6 +767,44 @@ class TestCircuitDrawerIntegration:
         )
         assert tape.draw() == expected
 
+    def test_qubit_circuit_with_max_length_kwdarg(self):
+        """Test that a qubit circuit with max_length set to 30 renders correctly."""
+        with qml.tape.QuantumTape() as tape:
+            for i in range(3):
+                qml.Hadamard(wires=i)
+                qml.RX(i * 0.1, wires=i)
+                qml.RY(i * 0.1, wires=i)
+                qml.RZ(i * 0.1, wires=i)
+            qml.expval(qml.PauliZ(0))
+
+        expected = (
+            " 0: ──H──RX(0)────RY(0)────RZ\n"
+            + " 1: ──H──RX(0.1)──RY(0.1)──RZ\n"
+            + " 2: ──H──RX(0.2)──RY(0.2)──RZ\n"
+            + "\n"
+            + " (0)────┤ ⟨Z⟩\n"
+            + " (0.1)──┤    \n"
+            + " (0.2)──┤    \n"
+        )
+        assert tape.draw(max_length=30) == expected
+
+    def test_qubit_circuit_length_under_max_length_kwdarg(self):
+        """Test that a qubit circuit with a circuit length less than the max_length renders correctly."""
+        with qml.tape.QuantumTape() as tape:
+            for i in range(3):
+                qml.Hadamard(wires=i)
+                qml.RX(i * 0.1, wires=i)
+                qml.RY(i * 0.1, wires=i)
+                qml.RZ(i * 0.1, wires=i)
+            qml.expval(qml.PauliZ(0))
+
+        expected = (
+            " 0: ──H──RX(0)────RY(0)────RZ(0)────┤ ⟨Z⟩\n"
+            + " 1: ──H──RX(0.1)──RY(0.1)──RZ(0.1)──┤    \n"
+            + " 2: ──H──RX(0.2)──RY(0.2)──RZ(0.2)──┤    \n"
+        )
+        assert tape.draw(max_length=60) == expected
+
 
 class TestWireOrdering:
     """Tests for wire ordering functionality"""
diff --git a/tests/transforms/test_draw.py b/tests/transforms/test_draw.py
index 22cf33de1..886c4cd5e 100644
--- a/tests/transforms/test_draw.py
+++ b/tests/transforms/test_draw.py
@@ -328,6 +328,79 @@ def test_same_wire_multiple_measurements_many_obs():
     assert qml.draw(qnode)(0.3, 0.2) == expected
 
 
+def test_qubit_circuit_with_max_length_kwdarg():
+    """Test that drawing a QNode with max_length=30 results in
+    correct circuit drawing."""
+    dev = qml.device("default.qubit", wires=3)
+
+    @qml.beta.qnode(dev)
+    def qnode():
+        for i in range(3):
+            qml.Hadamard(wires=i)
+            qml.RX(i * 0.1, wires=i)
+            qml.RY(i * 0.1, wires=i)
+            qml.RZ(i * 0.1, wires=i)
+        return qml.expval(qml.PauliZ(0))
+
+        expected = (
+            " 0: ──H──RX(0)────RY(0)────RZ\n"
+            + " 1: ──H──RX(0.1)──RY(0.1)──RZ\n"
+            + " 2: ──H──RX(0.2)──RY(0.2)──RZ\n"
+            + " \n"
+            + " (0)────┤ ⟨Z⟩ \n"
+            + " (0.1)──┤     \n"
+            + " (0.2)──┤     \n"
+        )
+        assert qml.draw(qnode, max_length=30)() == expected
+
+
+def test_qubit_circuit_with_max_length_kwdarg():
+    """Test that drawing a QNode with max_length=30 results in
+    correct circuit drawing."""
+    dev = qml.device("default.qubit", wires=3)
+
+    @qml.beta.qnode(dev)
+    def qnode():
+        for i in range(3):
+            qml.Hadamard(wires=i)
+            qml.RX(i * 0.1, wires=i)
+            qml.RY(i * 0.1, wires=i)
+            qml.RZ(i * 0.1, wires=i)
+        return qml.expval(qml.PauliZ(0))
+
+        expected = (
+            " 0: ──H──RX(0)────RY(0)────RZ\n"
+            + " 1: ──H──RX(0.1)──RY(0.1)──RZ\n"
+            + " 2: ──H──RX(0.2)──RY(0.2)──RZ\n"
+            + " \n"
+            + " (0)────┤ ⟨Z⟩ \n"
+            + " (0.1)──┤     \n"
+            + " (0.2)──┤     \n"
+        )
+        assert qml.draw(qnode, max_length=30)() == expected
+
+
+def test_qubit_circuit_length_under_max_length_kwdarg():
+    """Test that a qubit circuit with a circuit length less than the max_length renders correctly."""
+    dev = qml.device("default.qubit", wires=3)
+
+    @qml.beta.qnode(dev)
+    def qnode():
+        for i in range(3):
+            qml.Hadamard(wires=i)
+            qml.RX(i * 0.1, wires=i)
+            qml.RY(i * 0.1, wires=i)
+            qml.RZ(i * 0.1, wires=i)
+        return qml.expval(qml.PauliZ(0))
+
+        expected = (
+            " 0: ──H──RX(0)────RY(0)────RZ(0)────┤ ⟨Z⟩\n"
+            + " 1: ──H──RX(0.1)──RY(0.1)──RZ(0.1)──┤    \n"
+            + " 2: ──H──RX(0.2)──RY(0.2)──RZ(0.2)──┤    \n"
+        )
+        assert qml.draw(qnode, max_length=60)() == expected
+
+
 class TestWireOrdering:
     """Tests for wire ordering functionality"""