Merge pull request #178 from nemanjaASE/issue-169-incomplete-error-handling

Add error handling in report_chart.py

Author: msoedov

agentic_security/report_chart.py

@@ -7,11 +7,31 @@
 from matplotlib.cm import ScalarMappable
 from matplotlib.colors import LinearSegmentedColormap, Normalize
 
+from agentic_security.logutils import logger
+
 from .primitives import Table
 
 
 def plot_security_report(table: Table) -> io.BytesIO:
+    try:
+        return _plot_security_report(table=table)
+    except (TypeError, ValueError, OverflowError, IndexError, Exception) as e:
+        logger.error(f"Error in generating the security report: {e}")
+    return io.BytesIO()
+
+
+def generate_identifiers(data: pd.DataFrame) -> list[str]:
+    try:
+        _generate_identifiers(data=data)
+    except (TypeError, ValueError, Exception) as e:
+        logger.error(f"Error in generate_identifiers: {e}")
+    return [""]
+
+
+def _plot_security_report(table: Table) -> io.BytesIO:
     # Data preprocessing
+    logger.info("Data preprocessing started.")
+
     data = pd.DataFrame(table)
 
     # Sort by failure rate and reset index
@@ -22,10 +42,10 @@ def plot_security_report(table: Table) -> io.BytesIO:
     fig, ax = plt.subplots(figsize=(12, 10), subplot_kw={"projection": "polar"})
     fig.set_facecolor("#f0f0f0")
     ax.set_facecolor("#f0f0f0")
+    logger.info("Plot setup complete.")
 
     # Styling parameters
     colors = ["#6C5B7B", "#C06C84", "#F67280", "#F8B195"][::-1]  # Pastel palette
-    # colors = ["#440154", "#3b528b", "#21908c", "#5dc863"]  # Viridis-inspired palette
     cmap = LinearSegmentedColormap.from_list("custom", colors, N=256)
     norm = Normalize(vmin=data["tokens"].min(), vmax=data["tokens"].max())
 
@@ -76,7 +96,10 @@ def plot_security_report(table: Table) -> io.BytesIO:
 
     # Title and caption
     fig.suptitle(
-        "Security Report for Different Modules", fontsize=16, fontweight="bold", y=1.02
+        "Security Report for Different Modules",
+        fontsize=16,
+        fontweight="bold",
+        y=1.02,
     )
     caption = "Report generated by https://github.com/msoedov/agentic_security"
     fig.text(
@@ -114,17 +137,12 @@ def plot_security_report(table: Table) -> io.BytesIO:
             data["identifier"], data["failureRate"], data["module"]
         )
     ]
-    table = ax.table(
-        cellText=table_data,
-        loc="right",
-        cellLoc="left",
-    )
+    table = ax.table(cellText=table_data, loc="right", cellLoc="left")
     table.auto_set_font_size(False)
     table.set_fontsize(8)
 
     # Adjust table style
     table.scale(1, 0.7)
-
     for (row, col), cell in table.get_celld().items():
         cell.set_edgecolor("none")
         cell.set_facecolor("#f0f0f0" if row % 2 == 0 else "#e0e0e0")
@@ -134,17 +152,18 @@ def plot_security_report(table: Table) -> io.BytesIO:
             cell.set_text_props(fontweight="bold")
 
     # Adjust layout and save
-
     plt.tight_layout()
     buf = io.BytesIO()
     plt.savefig(buf, format="png", dpi=300, bbox_inches="tight")
     plt.close(fig)
     buf.seek(0)
+    logger.info("Report successfully generated and saved to buffer.")
     return buf
 
 
-def generate_identifiers(data: pd.DataFrame) -> list[str]:
+def _generate_identifiers(data: pd.DataFrame) -> list[str]:
     data_length = len(data)
+
     alphabet = string.ascii_uppercase
     num_letters = len(alphabet)