Update app.py

app.py

@@ -246,8 +246,11 @@ UNSAFE = re.compile(
     r"\bimport\s+(os|sys)\b|\bopen\s*\(|\beval\s*\(|\bexec\s*\(|\bgetattr\s*\()"
 )
 
-MAX_ANNOT_COLS = 12          # above this, heatmap cell numbers are unreadable
-MAX_TICKS = 20               # above this, thin the tick labels
+MAX_ANNOT_COLS = 12      # above this, heatmap cell numbers are unreadable
+MAX_TICKS = 20           # above this, thin the tick labels
+MAX_CAT_LEVELS = 20      # above this, a categorical plot has nothing to show
+SKIP_ATTR = "_edagent_skip"
+TITLE_ATTR = "_edagent_title"
 
 
 def sanitize(code: str) -> str:
@@ -292,10 +295,53 @@ def _tidy(fig):
 
 def _install_plot_guards():
     """Patch the real modules, so `import matplotlib.pyplot as plt` cannot bypass us."""
-    real = {"savefig": plt.savefig, "heatmap": sns.heatmap, "histplot": sns.histplot}
+    real = {"savefig": plt.savefig, "heatmap": sns.heatmap, "histplot": sns.histplot,
+            "boxplot": sns.boxplot, "violinplot": sns.violinplot,
+            "countplot": sns.countplot}
+
+    # ---- helpers ---------------------------------------------------------- #
+    def _vars(a, kw):
+        """Recover the (name, Series) pairs a seaborn call is plotting."""
+        d = kw.get("data", a[0] if a else None)
+        out = []
+        for key in ("x", "y"):
+            v = kw.get(key)
+            if isinstance(v, str) and hasattr(d, "columns"):
+                out.append((v, d[v]))
+            elif isinstance(v, pd.Series):
+                out.append((v.name, v))
+        if not out and isinstance(d, pd.Series):
+            out.append((d.name, d))
+        return out
+
+    def _counts_instead(name, s):
+        """A box plot of labels is meaningless; show how often each level occurs."""
+        ax = plt.gca()
+        s.value_counts().head(MAX_CAT_LEVELS).sort_values().plot(kind="barh", ax=ax)
+        ax.set_xlabel("count")
+        ax.set_ylabel(str(name or ""))
+        setattr(plt.gcf(), TITLE_ATTR, f"Counts of {name}" if name else "Counts")
+        return ax
 
+    def _categorical_guard(a, kw):
+        """Return an axes if we handled it, else None to fall through."""
+        pairs = _vars(a, kw)
+        if not pairs or any(pd.api.types.is_numeric_dtype(s) for _, s in pairs):
+            return None                       # a numeric variable is present: legitimate
+        name, s = pairs[0]
+        if s.nunique(dropna=True) > MAX_CAT_LEVELS:
+            setattr(plt.gcf(), SKIP_ATTR, True)   # company_name, free-text descriptions
+            return plt.gca()
+        return _counts_instead(name, s)
+
+    # ---- wrappers --------------------------------------------------------- #
     def savefig(fname, *a, **kw):
         fig = plt.gcf()
+        if getattr(fig, SKIP_ATTR, False):
+            return None                        # nothing worth writing to disk
+        title = getattr(fig, TITLE_ATTR, None)
+        if title and fig.get_axes():
+            fig.get_axes()[0].set_title(title)  # the model's "Box Plot of X" would lie
         w, h = fig.get_size_inches()
         fig.set_size_inches(max(w, 8), max(h, 5))
         _tidy(fig)
@@ -319,19 +365,15 @@ def _install_plot_guards():
                                rotation=0, fontsize=fs)
         return ax
 
-    def _series(a, kw):
-        d, x = kw.get("data", a[0] if a else None), kw.get("x")
-        try:
-            if isinstance(x, str) and hasattr(d, "columns"):
-                return pd.to_numeric(d[x], errors="coerce").dropna()
-            if isinstance(d, pd.Series):
-                return pd.to_numeric(d, errors="coerce").dropna()
-        except Exception:
-            pass
+    def _numeric_series(a, kw):
+        for _, s in _vars(a, kw):
+            s = pd.to_numeric(s, errors="coerce").dropna()
+            if len(s):
+                return s
         return None
 
     def histplot(*a, **kw):
-        s = _series(a, kw)
+        s = _numeric_series(a, kw)
         logged = False
         # heavy right skew (prices, fares, incomes) -> log x, else one tall bar
         if (s is not None and len(s) > 20 and s.min() > 0
@@ -347,7 +389,24 @@ def _install_plot_guards():
             ax.set_xlabel(f"{ax.get_xlabel()} (log scale)")
         return ax
 
-    plt.savefig, sns.heatmap, sns.histplot = savefig, heatmap, histplot
+    def boxplot(*a, **kw):
+        handled = _categorical_guard(a, kw)
+        return handled if handled is not None else real["boxplot"](*a, **kw)
+
+    def violinplot(*a, **kw):
+        handled = _categorical_guard(a, kw)
+        return handled if handled is not None else real["violinplot"](*a, **kw)
+
+    def countplot(*a, **kw):
+        pairs = _vars(a, kw)
+        if pairs and pairs[0][1].nunique(dropna=True) > MAX_CAT_LEVELS:
+            setattr(plt.gcf(), SKIP_ATTR, True)
+            return plt.gca()
+        return real["countplot"](*a, **kw)
+
+    plt.savefig = savefig
+    sns.heatmap, sns.histplot = heatmap, histplot
+    sns.boxplot, sns.violinplot, sns.countplot = boxplot, violinplot, countplot
     return real
 
 
@@ -367,13 +426,16 @@ def run_code(code: str, df: pd.DataFrame) -> dict:
     except Exception:
         err = traceback.format_exc(limit=3)
     finally:
+        # safety net: plt.show(), or a crash with a figure still open
         for i, num in enumerate(plt.get_fignums(), start=1):
             fig = plt.figure(num)
-            if fig.get_axes():
+            if fig.get_axes() and not getattr(fig, SKIP_ATTR, False):
                 fig.savefig(f"{PLOTS_DIR}/plot_{i}.png")
         plt.close("all")
-        plt.savefig, sns.heatmap, sns.histplot = (
-            real["savefig"], real["heatmap"], real["histplot"])
+        plt.savefig = real["savefig"]
+        sns.heatmap, sns.histplot = real["heatmap"], real["histplot"]
+        sns.boxplot, sns.violinplot = real["boxplot"], real["violinplot"]
+        sns.countplot = real["countplot"]
 
     return {"ok": err is None, "stdout": buf.getvalue(), "error": err,
             "plots": sorted(set(glob.glob(f"{PLOTS_DIR}/*.png")) - before)}