support for adding figures via figure handles

Author: headtr1ck

pylatex/figure.py

@@ -22,7 +22,7 @@ def add_image(
         filename,
         *,
         width=NoEscape(r"0.8\textwidth"),
-        placement=NoEscape(r"\centering")
+        placement=NoEscape(r"\centering"),
     ):
         """Add an image to the figure.
 
@@ -50,7 +50,7 @@ def add_image(
             StandAloneGraphic(image_options=width, filename=fix_filename(filename))
         )
 
-    def _save_plot(self, *args, extension="pdf", **kwargs):
+    def _save_plot(self, *args, figure=None, extension="pdf", **kwargs):
         """Save the plot.
 
         Returns
@@ -64,11 +64,12 @@ def _save_plot(self, *args, extension="pdf", **kwargs):
         filename = "{}.{}".format(str(uuid.uuid4()), extension.strip("."))
         filepath = posixpath.join(tmp_path, filename)
 
-        plt.savefig(filepath, *args, **kwargs)
+        fig = figure or plt.gcf()
+        fig.savefig(filepath, *args, **kwargs)
         return filepath
 
-    def add_plot(self, *args, extension="pdf", **kwargs):
-        """Add the current Matplotlib plot to the figure.
+    def add_plot(self, *args, figure=None, extension="pdf", **kwargs):
+        """Add a Matplotlib plot to the figure.
 
         The plot that gets added is the one that would normally be shown when
         using ``plt.show()``.
@@ -77,6 +78,8 @@ def add_plot(self, *args, extension="pdf", **kwargs):
         ----
         args:
             Arguments passed to plt.savefig for displaying the plot.
+        figure:
+            Optional matplotlib figure. If None add the current figure.
         extension : str
             extension of image file indicating figure file type
         kwargs:
@@ -92,7 +95,7 @@ def add_plot(self, *args, extension="pdf", **kwargs):
             if key in kwargs:
                 add_image_kwargs[key] = kwargs.pop(key)
 
-        filename = self._save_plot(*args, extension=extension, **kwargs)
+        filename = self._save_plot(*args, figure=figure, extension=extension, **kwargs)
 
         self.add_image(filename, **add_image_kwargs)
 

tests/test_pictures.py

@@ -4,16 +4,37 @@
 
 from pylatex import Document, Section
 from pylatex.figure import Figure
+import matplotlib.pyplot as plt
 
 
-def test():
+def test_add_image():
     doc = Document()
-    section = Section("Multirow Test")
+    section = Section("Add image Test")
     figure = Figure()
     image_filename = os.path.join(os.path.dirname(__file__), "../examples/kitten.jpg")
     figure.add_image(image_filename)
-    figure.add_caption("Whoooo an imagage of a pdf")
+    figure.add_caption("Whoooo an image of a kitty")
     section.append(figure)
     doc.append(section)
 
     doc.generate_pdf()
+
+
+def test_add_plot():
+    doc = Document()
+    section = Section("Add plot Test")
+    mplfig = plt.figure()
+
+    figure = Figure()
+    figure.add_plot()
+    figure.add_caption("Whoooo current matplotlib fig")
+    section.append(figure)
+
+    figure = Figure()
+    figure.add_plot(figure=mplfig)
+    figure.add_caption("Whoooo image from figure handle")
+    section.append(figure)
+
+    doc.append(section)
+
+    doc.generate_pdf()