support for adding figures via figure handles

Author: headtr1ck

pylatex/figure.py

@@ -44,7 +44,7 @@ def add_image(self, filename, *, width=NoEscape(r'0.8\textwidth'),
         self.append(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
@@ -58,11 +58,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()``.
@@ -71,6 +72,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:
@@ -86,7 +89,9 @@ 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

@@ -3,17 +3,38 @@
 from pylatex import Document, Section
 from pylatex.figure import Figure
 import os
+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()