feat(server): load_project_stat_klasses + project_root on /load_expr

buckaroo/server/handlers.py

@@ -344,9 +344,15 @@ async def post(self):
         no_browser = bool(body.get("no_browser", False))
         component_config = body.get("component_config")
 
+        project_root = body.get("project_root")
+
         try:
             expr = xorq_loading.load_expr_build_dir(build_dir)
-            xorq_dataflow = xorq_loading.XorqServerDataflow(expr, skip_main_serial=True)
+            extra_klasses = (
+                xorq_loading.load_project_stat_klasses(project_root)
+                if project_root else [])
+            xorq_dataflow = xorq_loading.XorqServerDataflow(
+                expr, skip_main_serial=True, extra_klasses=extra_klasses)
             metadata = xorq_loading.get_xorq_metadata(xorq_dataflow, build_dir)
         except FileNotFoundError:
             self.set_status(404)

buckaroo/server/xorq_loading.py

@@ -8,13 +8,19 @@
 """
 from __future__ import annotations
 
+import builtins
+import inspect
+import logging
 import traceback
+from pathlib import Path
 
 from buckaroo.xorq_buckaroo import (
     NoCleaningConfXorq, XorqAutocleaning, XorqDataflow, XorqDfStatsV2,
     XorqInfiniteSampling, _XORQ_ANALYSIS_KLASSES, _expr_count,
     window_to_parquet)
 
+log = logging.getLogger(__name__)
+
 
 class XorqServerDataflow(XorqDataflow):
     """Headless XorqDataflow with infinite sampling.
@@ -23,6 +29,13 @@ class XorqServerDataflow(XorqDataflow):
     its ``InnerDataFlow`` (xorq_buckaroo.py:181-186) — without
     ``InnerDataFlow``'s widget-side ``_df_to_obj`` override, since the
     server never serialises a main-frame sample (``skip_main_serial=True``).
+
+    ``extra_klasses`` is an optional list of additional ``@stat()``-decorated
+    functions (or ``ColAnalysis`` subclasses) to fold into ``analysis_klasses``
+    at the per-instance level. ``LoadExprHandler.post`` uses it to inject
+    project-authored stats discovered under ``<project_root>/stats/*.py``;
+    the built-in xorq stats are kept first so collisions resolve to the
+    built-in.
     """
 
     sampling_klass = XorqInfiniteSampling
@@ -31,6 +44,14 @@ class XorqServerDataflow(XorqDataflow):
     DFStatsClass = XorqDfStatsV2
     analysis_klasses = _XORQ_ANALYSIS_KLASSES
 
+    def __init__(self, expr, *args, extra_klasses=None, **kwargs):
+        if extra_klasses:
+            # Per-instance override — class-level _XORQ_ANALYSIS_KLASSES is
+            # left untouched so other sessions / direct widget usage don't
+            # inherit one project's stats.
+            self.analysis_klasses = list(_XORQ_ANALYSIS_KLASSES) + list(extra_klasses)
+        super().__init__(expr, *args, **kwargs)
+
 
 def load_expr_build_dir(build_dir: str):
     """Rehydrate an ibis expression from a xorq build directory.
@@ -59,6 +80,116 @@ def get_xorq_metadata(xorq_dataflow: XorqServerDataflow, build_dir: str) -> dict
     return {"path": build_dir, "rows": _expr_count(expr), "columns": columns}
 
 
+# ---------------------------------------------------------------------------
+# project-authored summary stats (loaded from <project_root>/stats/*.py)
+# ---------------------------------------------------------------------------
+
+# Names from ``builtins`` that the exec'd stat source is allowed to see.
+# Notable absences: ``__import__``, ``open``, ``exec``, ``eval``, ``compile``,
+# ``input``, ``getattr``/``setattr``/``delattr`` — without those the function
+# body cannot reach the filesystem, the import system, or the live process
+# graph through name lookup. Not a real sandbox (a determined caller can
+# walk ``col.__class__.__mro__``); acceptable because the project owner is
+# trusted (the buckaroo subprocess only reads project paths a host like
+# pydata-app explicitly passed in).
+_SAFE_BUILTIN_NAMES = ("True", "False", "None", "abs", "min", "max", "round", "sum", "len", "int", "float", "str", "bool", "list", "tuple", "dict", "set", "range", "enumerate", "zip", "map", "filter", "isinstance", "issubclass", "type")
+
+
+def _safe_builtins() -> dict:
+    return {n: getattr(builtins, n) for n in _SAFE_BUILTIN_NAMES}
+
+
+def load_project_stat_klasses(project_root) -> list:
+    """Scan ``<project_root>/stats/*.py`` and return wrapped ``@stat()`` funcs.
+
+    Each file must define a callable ``compute(col)`` returning an ibis
+    expression. The function is renamed to the filename stem and decorated
+    with ``@stat()`` from ``pluggable_analysis_framework.stat_func`` so it
+    slots into the same ``XORQ_STATS_V2`` list as the built-in xorq stats —
+    the ``XorqStatPipeline`` doesn't distinguish.
+
+    Files whose name starts with ``_`` (e.g. ``_disabled.py``, ``_helpers.py``)
+    are skipped, as a convention for parking work-in-progress without
+    removing it from the project directory.
+
+    Errors in any one file are logged and the file is skipped — one bad
+    stat shouldn't keep the rest from loading.
+    """
+    stats_dir = Path(project_root) / "stats"
+    if not stats_dir.is_dir():
+        return []
+
+    # Local imports — keep the noisy stat_func module surface off this
+    # file's import path. The @stat decorator + XorqColumn marker live
+    # in the pluggable framework rather than xorq itself.
+    from buckaroo.pluggable_analysis_framework.stat_func import (
+        XorqColumn, stat)
+
+    klasses: list = []
+    for path in sorted(stats_dir.glob("*.py")):
+        if path.name.startswith("_"):
+            continue
+        name = path.stem
+        if not name.isidentifier():
+            log.warning(
+                "project stat %s: filename stem is not a valid python identifier; skipping",
+                path)
+            continue
+        try:
+            wrapped = _compile_project_stat(name, path, stat, XorqColumn)
+        except Exception as e:
+            log.warning("project stat %s skipped: %s", path, e)
+            continue
+        klasses.append(wrapped)
+    return klasses
+
+
+def _compile_project_stat(name: str, path: Path, stat_decorator, XorqColumn):
+    """Read, exec, validate, and wrap one stat file. Raises on any failure;
+    the caller logs + skips. Kept separate so the per-file try/except in
+    the loader doesn't accidentally hide bugs in the loop bookkeeping."""
+    source = path.read_text()
+    globs: dict = {"__builtins__": _safe_builtins()}
+    # Expose ibis / xorq inside the function body for stats that need
+    # references beyond bare column methods (literal, case().when(), etc.).
+    try:
+        from xorq.vendor import ibis as _ibis  # noqa: PLC0415
+        globs["ibis"] = _ibis
+    except ImportError:
+        pass
+    try:
+        import xorq.api as _xo  # noqa: PLC0415
+        globs["xorq"] = _xo
+    except ImportError:
+        pass
+
+    # Single shared namespace for globals + locals so top-level
+    # constants and helper functions in the stat file are visible to
+    # ``compute`` when it runs. With separate dicts, ``compute`` captures
+    # ``globs`` as its ``__globals__`` while module-level names land in
+    # ``locals`` — every call would NameError.
+    exec(compile(source, str(path), "exec"), globs)
+
+    compute = globs.get("compute")
+    if not callable(compute):
+        raise ValueError("no callable 'compute' defined")
+
+    sig = inspect.signature(compute)
+    params = list(sig.parameters.values())
+    if len(params) != 1:
+        raise ValueError(
+            f"compute() must take exactly one parameter, got {len(params)}")
+
+    # Inject the XorqColumn annotation on the single param so the @stat
+    # decorator marks the function as needing raw column data (the
+    # pipeline then passes the actual ibis column expression in rather
+    # than looking up a stat keyed by the parameter name).
+    compute.__annotations__ = {params[0].name: XorqColumn}
+    compute.__name__ = name
+    compute.__qualname__ = name
+    return stat_decorator()(compute)
+
+
 def handle_infinite_request_xorq(xorq_dataflow: XorqServerDataflow,
         payload_args: dict) -> tuple[dict, bytes]:
     """Drive one infinite_request window against a xorq expression.

tests/unit/server/test_project_stats.py

@@ -0,0 +1,178 @@
+"""Tests for ``load_project_stat_klasses`` — scan a project's
+``stats/<name>.py`` directory, exec each file with restricted globals,
+return ``@stat()``-decorated callables that drop into the same
+``XORQ_STATS_V2`` list the built-in xorq stats live in.
+
+The feature lets a host (e.g. pydata-app's MCP server) hand buckaroo
+runtime-authored summary stats without restarting the server or adding
+new wire shapes. The function in each file is named ``compute`` and
+takes one positional argument (the column expression).
+"""
+from __future__ import annotations
+
+from pathlib import Path
+
+import pytest
+
+xo = pytest.importorskip("xorq.api")
+
+from buckaroo.server.xorq_loading import load_project_stat_klasses  # noqa: E402
+
+
+def test_returns_empty_when_stats_dir_missing(tmp_path: Path):
+    assert load_project_stat_klasses(tmp_path) == []
+
+
+def test_returns_empty_when_stats_dir_empty(tmp_path: Path):
+    (tmp_path / "stats").mkdir()
+    assert load_project_stat_klasses(tmp_path) == []
+
+
+def test_picks_up_one_stat_keyed_by_filename(tmp_path: Path):
+    stats = tmp_path / "stats"
+    stats.mkdir()
+    (stats / "percent_aa.py").write_text(
+        "def compute(col):\n"
+        "    return col.cast('string').contains('aa').sum() / col.count()\n")
+    klasses = load_project_stat_klasses(tmp_path)
+    assert len(klasses) == 1
+    # The wrapped function carries the filename stem as its stat name.
+    assert klasses[0]._stat_func.name == "percent_aa"
+
+
+def test_skips_file_without_compute(tmp_path: Path):
+    """A .py file that doesn't define compute() is skipped, not raised —
+    one bad file shouldn't block the rest of the stats from loading."""
+    stats = tmp_path / "stats"
+    stats.mkdir()
+    (stats / "no_compute.py").write_text("x = 42\n")
+    (stats / "n_rows.py").write_text("def compute(col): return col.count()\n")
+
+    klasses = load_project_stat_klasses(tmp_path)
+    names = sorted(k._stat_func.name for k in klasses)
+    assert names == ["n_rows"]
+
+
+def test_skips_file_that_tries_to_import(tmp_path: Path):
+    """Restricted globals strip __builtins__ so ``import os`` (which
+    resolves through ``__import__``) fails at exec. The file is logged
+    and skipped; the rest of the dir loads."""
+    stats = tmp_path / "stats"
+    stats.mkdir()
+    (stats / "evil.py").write_text(
+        "import os\n"
+        "def compute(col): return col.count()\n")
+    (stats / "n_rows.py").write_text("def compute(col): return col.count()\n")
+
+    klasses = load_project_stat_klasses(tmp_path)
+    names = sorted(k._stat_func.name for k in klasses)
+    assert names == ["n_rows"]
+
+
+def test_skips_underscore_prefixed_files(tmp_path: Path):
+    """Files starting with ``_`` (e.g. _disabled, _helpers) are not
+    treated as active stats — they let users park work-in-progress
+    without removing it from the project."""
+    stats = tmp_path / "stats"
+    stats.mkdir()
+    (stats / "_disabled.py").write_text("def compute(col): return col.count()\n")
+    (stats / "n_rows.py").write_text("def compute(col): return col.count()\n")
+
+    klasses = load_project_stat_klasses(tmp_path)
+    names = sorted(k._stat_func.name for k in klasses)
+    assert names == ["n_rows"]
+
+
+def test_loaded_stat_executes_against_an_ibis_column(tmp_path: Path):
+    """End-to-end: the wrapped function, called with a real ibis column,
+    returns an ibis expression that executes to the expected value.
+    This is the contract every stat in XORQ_STATS_V2 satisfies, and the
+    proof we haven't broken it by exec'ing the source out of a file."""
+    stats = tmp_path / "stats"
+    stats.mkdir()
+    (stats / "n_rows.py").write_text("def compute(col): return col.count()\n")
+
+    klasses = load_project_stat_klasses(tmp_path)
+    assert len(klasses) == 1
+
+    table = xo.memtable({"a": [1, 2, 3]}, name="t")
+    expr = klasses[0](table["a"])
+    assert int(expr.execute()) == 3
+
+
+def test_loaded_stat_can_use_module_level_constant(tmp_path: Path):
+    """A module-level constant defined in the stat file must be visible
+    to ``compute()`` when it runs. Codex P1: exec'ing with separate
+    globals / locals dicts puts top-level assignments into locals, but
+    the function captures globals as its ``__globals__`` — so the
+    constant is invisible at call time and the function NameErrors on
+    every column."""
+    stats = tmp_path / "stats"
+    stats.mkdir()
+    (stats / "above_threshold.py").write_text(
+        "THRESHOLD = 2\n"
+        "def compute(col): return (col > THRESHOLD).sum()\n")
+
+    klasses = load_project_stat_klasses(tmp_path)
+    assert len(klasses) == 1
+
+    table = xo.memtable({"a": [1, 2, 3]}, name="t")
+    expr = klasses[0](table["a"])
+    assert int(expr.execute()) == 1
+
+
+def test_loaded_stat_can_use_module_level_helper(tmp_path: Path):
+    """Same Codex P1 in helper-function form: a top-level ``def`` in
+    the stat file must be callable from ``compute()``."""
+    stats = tmp_path / "stats"
+    stats.mkdir()
+    (stats / "double_count.py").write_text(
+        "def _double(x):\n"
+        "    return x * 2\n"
+        "def compute(col):\n"
+        "    return _double(col.count())\n")
+
+    klasses = load_project_stat_klasses(tmp_path)
+    assert len(klasses) == 1
+
+    table = xo.memtable({"a": [1, 2, 3]}, name="t")
+    expr = klasses[0](table["a"])
+    assert int(expr.execute()) == 6
+
+
+def test_dataflow_extra_klasses_extends_analysis_klasses(tmp_path: Path):
+    """XorqServerDataflow's per-instance extra_klasses appends to
+    _XORQ_ANALYSIS_KLASSES without mutating the class-level list.
+    This is the surface the /load_expr handler uses to inject project
+    stats — the test pins it so handler-only changes can't silently
+    bypass the extension path."""
+    from buckaroo.server.xorq_loading import (
+        XorqServerDataflow, _XORQ_ANALYSIS_KLASSES)
+
+    stats = tmp_path / "stats"
+    stats.mkdir()
+    (stats / "n_rows.py").write_text("def compute(col): return col.count()\n")
+    extra = load_project_stat_klasses(tmp_path)
+    assert len(extra) == 1
+
+    expr = xo.memtable({"a": [1, 2, 3]}, name="t")
+    dataflow = XorqServerDataflow(expr, skip_main_serial=True, extra_klasses=extra)
+
+    # Instance attr extends; class attr untouched.
+    assert dataflow.analysis_klasses is not _XORQ_ANALYSIS_KLASSES
+    assert len(dataflow.analysis_klasses) == len(_XORQ_ANALYSIS_KLASSES) + 1
+    assert extra[0] in dataflow.analysis_klasses
+    # All built-ins still there.
+    for builtin in _XORQ_ANALYSIS_KLASSES:
+        assert builtin in dataflow.analysis_klasses
+
+
+def test_dataflow_without_extra_klasses_keeps_class_default(tmp_path: Path):
+    """No extra_klasses → instance falls back to the class-level
+    _XORQ_ANALYSIS_KLASSES (the existing behaviour, must not regress)."""
+    from buckaroo.server.xorq_loading import (
+        XorqServerDataflow, _XORQ_ANALYSIS_KLASSES)
+
+    expr = xo.memtable({"a": [1, 2, 3]}, name="t")
+    dataflow = XorqServerDataflow(expr, skip_main_serial=True)
+    assert dataflow.analysis_klasses is _XORQ_ANALYSIS_KLASSES