Update report.py

report.py

@@ -1,9 +1,8 @@
-# report.py
-
 import io
 from reportlab.lib.pagesizes import letter
 from reportlab.platypus import SimpleDocTemplate, Image, Paragraph, Spacer
 from reportlab.lib.styles import getSampleStyleSheet
+import matplotlib.pyplot as plt
 
 class ReportGenerator:
     def __init__(self):
@@ -33,7 +32,15 @@ class ReportGenerator:
         elements = []
         elements.append(Paragraph(title, self.styles['Heading2']))
         img_buffer = io.BytesIO()
-        fig.write_image(img_buffer, format="png")
+        
+        if hasattr(fig, 'write_image'):  # Plotly figure
+            fig.write_image(img_buffer, format="png")
+        elif isinstance(fig, plt.Figure):  # Matplotlib figure
+            fig.savefig(img_buffer, format='png')
+            plt.close(fig)  # Close the figure to free up memory
+        else:
+            raise ValueError(f"Unsupported figure type: {type(fig)}")
+        
         img_buffer.seek(0)
         elements.append(Image(img_buffer, width=500, height=300))
         elements.append(Spacer(1, 12))