[python] Allow kernel arguments with cudaq.draw's format overload

python/cudaq/runtime/draw.py

@@ -39,7 +39,7 @@ def draw(decoratorOrFormat, *args):
     """
     if isinstance(decoratorOrFormat, str):
         # User specified output format.
-        assert (len(args) == 1) and "must have a kernel"
+        assert len(args) >= 1, "must have a kernel"
         vargs = args[1:]
         return _detail_draw(decoratorOrFormat, args[0], *vargs)
     # Default to the UTF-8 code points (confusingly named `"ascii"`).

python/tests/visualization/test_draw.py

@@ -105,6 +105,26 @@ def kernel():
     assert expected_str == produced_string
 
 
+def test_draw_format_with_arguments():
+    """The `cudaq.draw("<format>", kernel, ...)` overload must forward the
+    kernel arguments, just like the `cudaq.draw(kernel, ...)` overload."""
+
+    @cudaq.kernel
+    def kernel(theta: float):
+        q = cudaq.qvector(2)
+        ry(theta, q[0])
+        x.ctrl(q[0], q[1])
+
+    produced_string = cudaq.draw("ascii", kernel, 0.59)
+    assert produced_string == cudaq.draw(kernel, 0.59)
+    assert "ry(0.59)" in produced_string
+    assert "R_y(0.59)" in cudaq.draw("latex", kernel, 0.59)
+
+    # A format string must still be followed by a kernel.
+    with pytest.raises(AssertionError, match="must have a kernel"):
+        cudaq.draw("ascii")
+
+
 # This test will run on the default simulator. For machines with GPUs, that
 # will be a GPU-accelerated simulator, but for machines without GPUs, it
 # will run on a CPU simulator.