diff --git a/python/cudaq/runtime/draw.py b/python/cudaq/runtime/draw.py index ede3e1ec0e8..fc4ec9ae8c7 100644 --- a/python/cudaq/runtime/draw.py +++ b/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"`). diff --git a/python/tests/visualization/test_draw.py b/python/tests/visualization/test_draw.py index 0798b08ee58..af239864a44 100644 --- a/python/tests/visualization/test_draw.py +++ b/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("", 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.