Operations/measurements on all wires are drawn correctly with mid-cir…

…cuit measurements (#5501)

**Context:**
Bug fix for #5499 

**Description of the Change:**
Adding labels to layer strings only extends to wires, not classical
bits.

**Benefits:**
Drawer works correctly with mid-circuit measurements

**Possible Drawbacks:**

**Related GitHub Issues:**

doc/releases/changelog-dev.md

@@ -339,6 +339,9 @@
 
 <h3>Bug fixes 🐛</h3>
 
+* Operators applied to all wires are now drawn correctly in a circuit with mid-circuit measurements.
+ [(#5501)](https://github.com/PennyLaneAI/pennylane/pull/5501)
+
 * Fix a bug where certain unary mid-circuit measurement expressions would raise an uncaught error.
  [(#5480)](https://github.com/PennyLaneAI/pennylane/pull/5480)
 

pennylane/drawer/tape_text.py

@@ -142,7 +142,8 @@ def _add_op(op, layer_str, config):
 
  label = op.label(decimals=config.decimals, cache=config.cache).replace("\n", "")
  if len(op.wires) == 0: # operation (e.g. barrier, snapshot) across all wires
- for i, s in enumerate(layer_str):
+ n_wires = len(config.wire_map)
+ for i, s in enumerate(layer_str[:n_wires]):
  layer_str[i] = s + label
  else:
  for w in op.wires:
@@ -225,7 +226,8 @@ def _add_measurement(m, layer_str, config):
  meas_label = m.return_type.value
 
  if len(m.wires) == 0: # state or probability across all wires
- for i, s in enumerate(layer_str):
+ n_wires = len(config.wire_map)
+ for i, s in enumerate(layer_str[:n_wires]):
  layer_str[i] = s + meas_label
 
  for w in m.wires:

tests/drawer/test_draw.py

@@ -284,6 +284,8 @@ def circ():
 class TestMidCircuitMeasurements:
  """Tests for drawing mid-circuit measurements and classical conditions."""
 
+ # pylint: disable=too-many-public-methods
+
  @pytest.mark.parametrize("device_name", ["default.qubit"])
  def test_qnode_mid_circuit_measurement_not_deferred(self, device_name, mocker):
  """Test that a circuit containing mid-circuit measurements is transformed by the drawer
@@ -329,6 +331,51 @@ def func():
 
  assert drawing == expected_drawing
 
+ @pytest.mark.parametrize(
+ "op", [qml.GlobalPhase(0.1), qml.Identity(), qml.Snapshot(), qml.Barrier()]
+ )
+ @pytest.mark.parametrize("decimals", [None, 2])
+ def test_draw_all_wire_ops(self, op, decimals):
+ """Test that operators acting on all wires are drawn correctly"""
+
+ def func():
+ qml.X(0)
+ qml.X(1)
+ m = qml.measure(0)
+ qml.cond(m, qml.X)(0)
+ qml.apply(op)
+ return qml.expval(qml.Z(0))
+
+ # Stripping to remove trailing white-space because length of white-space at the
+ # end of the drawing depends on the length of each individual line
+ drawing = qml.draw(func, decimals=decimals)().strip()
+ label = op.label(decimals=decimals).replace("\n", "")
+ expected_drawing = (
+ f"0: ──X──┤↗├──X──{label}─┤ <Z>\n1: ──X───║───║──{label}─┤ \n ╚═══╝"
+ )
+
+ assert drawing == expected_drawing
+
+ @pytest.mark.parametrize(
+ "mp, label", [(qml.sample(), "Sample"), (qml.probs(), "Probs"), (qml.counts(), "Counts")]
+ )
+ def test_draw_all_wire_measurements(self, mp, label):
+ """Test that operators acting on all wires are drawn correctly"""
+
+ def func():
+ qml.X(0)
+ qml.X(1)
+ m = qml.measure(0)
+ qml.cond(m, qml.X)(0)
+ return qml.apply(mp)
+
+ # Stripping to remove trailing white-space because length of white-space at the
+ # end of the drawing depends on the length of each individual line
+ drawing = qml.draw(func)().strip()
+ expected_drawing = f"0: ──X──┤↗├──X─┤ {label}\n1: ──X───║───║─┤ {label}\n ╚═══╝"
+
+ assert drawing == expected_drawing
+
  def test_draw_mid_circuit_measurement_multiple_wires(self):
  """Test that mid-circuit measurements are correctly drawn in circuits
  with multiple wires."""