Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions src/tallyman_core/summary_stats.py
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,16 @@
"isinstance",
"issubclass",
"type",
# Exception types — needed for try/except Exception: in compute()
"Exception",
"ValueError",
"TypeError",
"KeyError",
"AttributeError",
"IndexError",
"RuntimeError",
"StopIteration",
"NotImplementedError",
)
_NAME_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")

Expand Down
45 changes: 25 additions & 20 deletions src/tallyman_mcp/server.py
Original file line number Diff line number Diff line change
Expand Up @@ -203,31 +203,36 @@ def catalog_run(code: str, prompt: str = "") -> dict:
"""Execute a xorq expression script and persist it as a catalog entry.

The provided code MUST bind a top-level variable named `expr` to a xorq/ibis
expression. Imports allowed; all imports happen in a fresh module scope.
expression. The four canonical names are PRE-INJECTED — no import statement
needed for basic usage:

PREFERRED pattern:
ibis # xorq.vendor.ibis — ibis._, ibis.cases, ibis.window,
# ibis.literal, ibis.coalesce, ibis.desc, ibis.schema
xo # xorq.api — xo.memtable, xo.deferred_read_parquet, xo.connect
from_catalog # read a named catalog entry or content hash
from_project # read a raw file under <project>/data/

from tallyman_xorq.io import from_project, from_catalog
import xorq.vendor.ibis as ibis
t = from_catalog("alias_name") # named catalog entry (alias or hash)
MINIMAL PATTERN (no imports needed):

t = from_catalog("alias_name") # named entry
t = from_project("file.parquet") # raw file under <project>/data/
expr = t.filter(t.col > 0).group_by("region").aggregate(n=t.count())
expr = t.mutate(pk=ibis._.a + "_" + ibis._.b) # ibis._ available pre-injected

Explicit `import xorq.api as xo` or `import xorq.vendor.ibis as ibis` still work
and are fine for clarity — they just overwrite the pre-injected binding with the
same module.

FORBIDDEN IMPORTS — the linter rejects these before any code runs:

import ibis # imports system ibis (no datafusion backend); use
# the pre-injected ibis or `import xorq.vendor.ibis as ibis`
from ibis import X # same problem — use xorq.vendor.ibis
import xorq as xo # bare xorq has no API methods; use `import xorq.api as xo`
import xorq # same — use `import xorq.api as xo`

THREE NAMESPACES — mixing these up is the #1 build failure:

import xorq.api as xo # backends + deferred reads:
# xo.memtable, xo.deferred_read_parquet, xo.connect
import xorq.vendor.ibis as ibis # the expression API: ibis._, ibis.cases,
# ibis.window, ibis.literal, ibis.coalesce, ibis.desc
from tallyman_xorq.io import from_project, from_catalog # reading data

t.mutate(pk=ibis._.a + "_" + ibis._.b) # deferred string concat via ibis._
# NEVER `import xorq` then `xorq.<fn>`: bare xorq.read_parquet / xorq.memtable /
# xorq._ all raise AttributeError. The api module is `xo` (xorq.api).
# NEVER bare `import ibis` / `from ibis ...`: it builds but fails at save time
# with a vendored-Expr class error. Always `import xorq.vendor.ibis as ibis`.
# Math is a COLUMN METHOD: col.sin(), col.log(), col.sqrt() — not ibis.sin(col).
# Read data only via from_project / from_catalog (no xo.read_parquet / ibis.read_parquet).
Math is a COLUMN METHOD: col.sin(), col.log(), col.sqrt() — not ibis.sin(col).
Read data only via from_project / from_catalog (no xo.read_parquet / ibis.read_parquet).

ONLY BACKEND — xorq's built-in datafusion; there is NO duckdb. Do not use
`ibis.duckdb`, a duckdb connection, `.sql()`, or `con.register()`. Build
Expand Down
2 changes: 2 additions & 0 deletions src/tallyman_xorq/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
BuildError,
BuildResult,
build_and_persist,
lint_catalog_code,
list_entries,
load_entry,
read_prompts,
Expand Down Expand Up @@ -47,6 +48,7 @@
"key_diff",
"key_diff_polars",
"key_diff_xorq",
"lint_catalog_code",
"list_entries",
"load_entry",
"project_path",
Expand Down
74 changes: 74 additions & 0 deletions src/tallyman_xorq/build.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
from __future__ import annotations

import ast
import importlib.util
import json
import re
Expand Down Expand Up @@ -73,6 +74,51 @@ def _user_imports_bare_ibis(code: str) -> bool:
return False


def lint_catalog_code(code: str) -> list[str]:
"""AST-based pre-execution scan. Returns actionable error strings, empty = clean.

Catches import mistakes that cause cryptic runtime errors. Runs before any
code is executed so the model gets a clear rejection on the first bad call
rather than after exec blows up with an AttributeError or IbisError.
"""
try:
tree = ast.parse(code)
except SyntaxError as e:
return [f"syntax error: {e}"]

errors: list[str] = []
for node in ast.walk(tree):
if isinstance(node, ast.Import):
for alias in node.names:
if alias.name == "ibis":
errors.append(
"`import ibis` imports the system ibis package which has no datafusion "
"backend here and will shadow the pre-injected xorq.vendor.ibis. "
"ibis is already available without any import — just use `ibis.window(...)`, "
"`ibis.literal(...)`, `ibis._` etc. directly. If you need an explicit import "
"for clarity, write `import xorq.vendor.ibis as ibis`."
)
elif alias.name == "xorq":
stmt = f"`import xorq as {alias.asname}`" if alias.asname else "`import xorq`"
errors.append(
f"{stmt} — the bare `xorq` module has no API methods "
"(connect, memtable, deferred_read_parquet live on `xorq.api`). "
"Use `import xorq.api as xo`, or just use the pre-injected `xo` "
"directly without any import statement."
)
elif isinstance(node, ast.ImportFrom):
module = node.module or ""
if module == "ibis" or module.startswith("ibis."):
names = ", ".join(a.name for a in node.names)
errors.append(
f"`from {module} import {names}` imports from the system ibis package. "
"Use `import xorq.vendor.ibis as ibis` instead, or just use the pre-injected "
"`ibis` without any import statement."
)

return errors


# Read helpers the model reaches for on the wrong namespace. The project-aware
# way is from_project/from_catalog; xo.deferred_read_parquet handles a raw path.
_READ_FNS = frozenset(
Expand Down Expand Up @@ -210,6 +256,27 @@ def _ibis_import_hint(exc_msg: str, code: str = "") -> str:
return "\n\nHint: " + " ".join(hints)


def _inject_prelude(module: object) -> None:
"""Pre-inject the four canonical names into a module before exec_module runs.

ibis, xo, from_catalog, from_project are available without any import
statement. Explicit `import xorq.api as xo` / `import xorq.vendor.ibis as ibis`
still work — they just overwrite the same module object.
"""
try:
import xorq.api as _xo # noqa: PLC0415
import xorq.vendor.ibis as _ibis # noqa: PLC0415

from tallyman_xorq.io import from_catalog as _from_catalog # noqa: PLC0415
from tallyman_xorq.io import from_project as _from_project # noqa: PLC0415
except ImportError:
return
module.ibis = _ibis # type: ignore[attr-defined]
module.xo = _xo # type: ignore[attr-defined]
module.from_catalog = _from_catalog # type: ignore[attr-defined]
module.from_project = _from_project # type: ignore[attr-defined]


def _import_script(code: str) -> tuple[object, Path]:
"""Write code to a temp file, import it, return (module, temp_path).

Expand All @@ -221,6 +288,7 @@ def _import_script(code: str) -> tuple[object, Path]:
if spec is None or spec.loader is None:
raise BuildError(f"Could not load module spec from {tmp}")
module = importlib.util.module_from_spec(spec)
_inject_prelude(module)
sys.modules[spec.name] = module
try:
spec.loader.exec_module(module)
Expand All @@ -241,6 +309,12 @@ def build_and_persist(
The user code must bind a variable named `expr_name` (default "expr") to an
ibis/xorq expression. Imports happen in a fresh module scope.
"""
lint_errors = lint_catalog_code(code)
if lint_errors:
raise BuildError(
"pre-execution lint:\n" + "\n".join(f" - {e}" for e in lint_errors)
)

from xorq.ibis_yaml.compiler import build_expr, load_expr

from tallyman_core.paths import compute_cache_dir
Expand Down
82 changes: 81 additions & 1 deletion tests/test_build.py
Original file line number Diff line number Diff line change
Expand Up @@ -105,7 +105,87 @@ def test_build_no_hint_on_unrelated_error(project: str, orders_parquet: Path):
# hint only caught bare `import ibis`, the Expr class mismatch, and `xorq._`.
# ---------------------------------------------------------------------------

from tallyman_xorq.build import _ibis_import_hint # noqa: E402
from tallyman_xorq.build import _ibis_import_hint, lint_catalog_code # noqa: E402

# ---------------------------------------------------------------------------
# lint_catalog_code — pre-execution static scanner
# ---------------------------------------------------------------------------


def test_lint_rejects_import_ibis():
errors = lint_catalog_code("import ibis\nexpr = ibis.table('x')")
assert errors
assert any("xorq.vendor.ibis" in e for e in errors)


def test_lint_rejects_from_ibis_import():
errors = lint_catalog_code("from ibis import window\nexpr = None")
assert errors
assert any("xorq.vendor.ibis" in e for e in errors)


def test_lint_rejects_from_ibis_submodule():
errors = lint_catalog_code("from ibis.expr.types import Column\nexpr = None")
assert errors


def test_lint_rejects_bare_import_xorq():
errors = lint_catalog_code("import xorq\nexpr = xorq.connect()")
assert errors
assert any("xorq.api" in e for e in errors)


def test_lint_rejects_import_xorq_as_xo():
errors = lint_catalog_code("import xorq as xo\nexpr = xo.connect()")
assert errors
assert any("xorq.api" in e for e in errors)


def test_lint_accepts_correct_imports():
code = """
import xorq.api as xo
import xorq.vendor.ibis as ibis
from tallyman_xorq.io import from_catalog, from_project
t = from_catalog("foo")
expr = t.filter(t.x > 0)
"""
assert lint_catalog_code(code) == []


def test_lint_accepts_no_imports():
code = "t = from_catalog('foo')\nexpr = t.filter(t.x > 0)"
assert lint_catalog_code(code) == []


def test_lint_returns_syntax_error():
errors = lint_catalog_code("def broken(\nexpr = None")
assert errors
assert any("syntax" in e for e in errors)


def test_build_lint_fires_before_exec(project: str):
# `import ibis` is caught by the linter — no code should execute.
code = "import ibis\nexpr = ibis.table('x', schema={'a': 'int64'})"
with pytest.raises(BuildError, match="pre-execution lint"):
build_and_persist(project, code)


def test_preinjected_names_available(project: str, orders_parquet: Path):
# ibis, xo, from_project are pre-injected — no import statements needed.
code = f"""
t = xo.deferred_read_parquet({str(orders_parquet)!r})
expr = t.group_by("region").aggregate(
total=t.price.sum(),
label=ibis.literal("ok"),
)
"""
res = build_and_persist(project, code)
assert res.row_count == 4


# ---------------------------------------------------------------------------
# Namespace / column-name hints
# ---------------------------------------------------------------------------


def test_hint_bare_xorq_attribute_points_to_api():
Expand Down
Loading