Add saving possibility to register layout draw method (#446)

* Correcting typos

* Add fig-name and kwargs-savefig

* Delete blankspace

pulser-core/pulser/register/register_layout.py

@@ -197,6 +197,8 @@ def draw(
         draw_graph: bool = False,
         draw_half_radius: bool = False,
         projection: bool = True,
+        fig_name: str = None,
+        kwargs_savefig: dict = {},
     ) -> None:
         """Draws the entire register layout.
 
@@ -211,6 +213,11 @@ def draw(
                 if the `blockade_radius` is defined.
             projection: If the layout is in 3D, draws it
                 as projections on different planes.
+            fig_name: The name on which to save the figure.
+                If None the figure will not be saved.
+            kwargs_savefig: Keywords arguments for
+                ``matplotlib.pyplot.savefig``. Not applicable if `fig_name`
+                is ``None``.
 
         Note:
             When drawing half the blockade radius, we say there is a blockade
@@ -252,6 +259,8 @@ def draw(
                 draw_half_radius=draw_half_radius,
                 are_traps=True,
             )
+        if fig_name is not None:
+            plt.savefig(fig_name, **kwargs_savefig)
         plt.show()
 
     def make_mappable_register(

tests/test_register_layout.py

@@ -114,12 +114,20 @@ def test_draw(layout, layout3d):
     with patch("matplotlib.pyplot.show"):
         layout.draw()
 
+    with patch("matplotlib.pyplot.show"):
+        with patch("matplotlib.pyplot.savefig"):
+            layout.draw(fig_name="my_registerlayout.pdf")
+
     with patch("matplotlib.pyplot.show"):
         layout3d.draw()
 
     with patch("matplotlib.pyplot.show"):
         layout3d.draw(projection=False)
 
+    with patch("matplotlib.pyplot.show"):
+        with patch("matplotlib.pyplot.savefig"):
+            layout3d.draw(fig_name="my_registerlayout.pdf")
+
 
 def test_repr(layout):
     hash_ = sha256(bytes(2))