From adbe88707cba6ee327e0663563b8dd1713c4dfc6 Mon Sep 17 00:00:00 2001
From: Matthew Treinish <mtreinish@kortar.org>
Date: Wed, 7 Aug 2024 20:35:20 -0400
Subject: [PATCH] Fix dag visualization with Var wires (#12848)

* Fix dag visualization with Var wires

This commit fixes the dag visualization for DAGs with classical
variables. The Var type was not handled in the attribute callback
functions for nodes and edges. This was causing the visualizer to fail
if the dag contained these types. This fixes this by adding explict
handling for the Var types and using the name attribute of the Var
object.

* Add release note and test
---
 qiskit/visualization/dag_visualization.py         | 13 ++++++++++---
 .../notes/fix-var-wires-4ebc40e0b19df253.yaml     |  8 ++++++++
 test/python/visualization/test_dag_drawer.py      | 15 ++++++++++++++-
 3 files changed, 32 insertions(+), 4 deletions(-)
 create mode 100644 releasenotes/notes/fix-var-wires-4ebc40e0b19df253.yaml

diff --git a/qiskit/visualization/dag_visualization.py b/qiskit/visualization/dag_visualization.py
index ad2fca6e9b..dae97b51b0 100644
--- a/qiskit/visualization/dag_visualization.py
+++ b/qiskit/visualization/dag_visualization.py
@@ -174,10 +174,13 @@ def dag_drawer(dag, scale=0.7, filename=None, style="color"):
                         label = register_bit_labels.get(
                             node.wire, f"q_{dag.find_bit(node.wire).index}"
                         )
-                    else:
+                    elif isinstance(node.wire, Clbit):
                         label = register_bit_labels.get(
                             node.wire, f"c_{dag.find_bit(node.wire).index}"
                         )
+                    else:
+                        label = str(node.wire.name)
+
                     n["label"] = label
                     n["color"] = "black"
                     n["style"] = "filled"
@@ -187,10 +190,12 @@ def dag_drawer(dag, scale=0.7, filename=None, style="color"):
                         label = register_bit_labels.get(
                             node.wire, f"q[{dag.find_bit(node.wire).index}]"
                         )
-                    else:
+                    elif isinstance(node.wire, Clbit):
                         label = register_bit_labels.get(
                             node.wire, f"c[{dag.find_bit(node.wire).index}]"
                         )
+                    else:
+                        label = str(node.wire.name)
                     n["label"] = label
                     n["color"] = "black"
                     n["style"] = "filled"
@@ -203,8 +208,10 @@ def dag_drawer(dag, scale=0.7, filename=None, style="color"):
             e = {}
             if isinstance(edge, Qubit):
                 label = register_bit_labels.get(edge, f"q_{dag.find_bit(edge).index}")
-            else:
+            elif isinstance(edge, Clbit):
                 label = register_bit_labels.get(edge, f"c_{dag.find_bit(edge).index}")
+            else:
+                label = str(edge.name)
             e["label"] = label
             return e
 
diff --git a/releasenotes/notes/fix-var-wires-4ebc40e0b19df253.yaml b/releasenotes/notes/fix-var-wires-4ebc40e0b19df253.yaml
new file mode 100644
index 0000000000..7cd1e74806
--- /dev/null
+++ b/releasenotes/notes/fix-var-wires-4ebc40e0b19df253.yaml
@@ -0,0 +1,8 @@
+---
+fixes:
+  - |
+    Fixed an issue with :func:`.dag_drawer` and :meth:`.DAGCircuit.draw`
+    when attempting to visualize a :class:`.DAGCircuit` instance that contained
+    :class:`.Var` wires. The visualizer would raise an exception trying to
+    do this which has been fixed so the expected visualization will be
+    generated.
diff --git a/test/python/visualization/test_dag_drawer.py b/test/python/visualization/test_dag_drawer.py
index 4b920390e8..d789b1e706 100644
--- a/test/python/visualization/test_dag_drawer.py
+++ b/test/python/visualization/test_dag_drawer.py
@@ -16,12 +16,14 @@ import os
 import tempfile
 import unittest
 
-from qiskit.circuit import QuantumRegister, ClassicalRegister, QuantumCircuit, Qubit, Clbit
+from qiskit.circuit import QuantumRegister, ClassicalRegister, QuantumCircuit, Qubit, Clbit, Store
 from qiskit.visualization import dag_drawer
 from qiskit.exceptions import InvalidFileError
 from qiskit.visualization import VisualizationError
 from qiskit.converters import circuit_to_dag, circuit_to_dagdependency
 from qiskit.utils import optionals as _optionals
+from qiskit.dagcircuit import DAGCircuit
+from qiskit.circuit.classical import expr, types
 from .visualization import path_to_diagram_reference, QiskitVisualizationTestCase
 
 
@@ -108,6 +110,17 @@ class TestDagDrawer(QiskitVisualizationTestCase):
             image = Image.open(tmp_path)
             self.assertImagesAreEqual(image, image_ref, 0.1)
 
+    @unittest.skipUnless(_optionals.HAS_GRAPHVIZ, "Graphviz not installed")
+    @unittest.skipUnless(_optionals.HAS_PIL, "PIL not installed")
+    def test_dag_drawer_with_var_wires(self):
+        """Test visualization works with var nodes."""
+        a = expr.Var.new("a", types.Bool())
+        dag = DAGCircuit()
+        dag.add_input_var(a)
+        dag.apply_operation_back(Store(a, a), (), ())
+        image = dag_drawer(dag)
+        self.assertIsNotNone(image)
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)