--functions can take filenames, can be used multiple times (#681)

Closes #659

The --functions option now accepts:
- File paths ending in .py (e.g., --functions my_funcs.py)
- Multiple invocations (e.g., --functions foo.py --functions 'def bar(): ...')
- Inline Python code (existing behavior)

Implementation follows the same pattern as llm's --functions flag
(simonw/llm@a880c123).

Changes:
- Added multiple=True to --functions Click option in query, bulk, and memory commands
- Modified _register_functions() to detect and read .py files
- Updated _maybe_register_functions() to iterate over multiple function sources
- Removed unused bytes/bytearray handling
- Added comprehensive tests for file paths and multiple invocations
- Updated documentation with examples

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>

* Shorter help for --functions

---------

Co-authored-by: Claude <noreply@anthropic.com>
This commit is contained in:
Simon Willison 2025-11-23 21:46:51 -08:00
commit 638b8ca871
6 changed files with 166 additions and 9 deletions

View file

@ -962,7 +962,7 @@ def insert_upsert_implementation(
db = sqlite_utils.Database(path)
_load_extensions(db, load_extension)
if functions:
_register_functions(db, functions)
_maybe_register_functions(db, functions)
if (delimiter or quotechar or sniff or no_headers) and not tsv:
csv = True
if (nl + csv + tsv) >= 2:
@ -1370,7 +1370,9 @@ def upsert(
@click.argument("file", type=click.File("rb"), required=True)
@click.option("--batch-size", type=int, default=100, help="Commit every X records")
@click.option(
"--functions", help="Python code defining one or more custom SQL functions"
"--functions",
help="Python code or file path defining custom SQL functions",
multiple=True,
)
@import_options
@load_extension_option
@ -1759,7 +1761,9 @@ def drop_view(path, view, ignore, load_extension):
help="Named :parameters for SQL query",
)
@click.option(
"--functions", help="Python code defining one or more custom SQL functions"
"--functions",
help="Python code or file path defining custom SQL functions",
multiple=True,
)
@load_extension_option
def query(
@ -1796,7 +1800,7 @@ def query(
db.register_fts4_bm25()
if functions:
_register_functions(db, functions)
_maybe_register_functions(db, functions)
_execute_query(
db,
@ -1824,7 +1828,9 @@ def query(
)
@click.argument("sql")
@click.option(
"--functions", help="Python code defining one or more custom SQL functions"
"--functions",
help="Python code or file path defining custom SQL functions",
multiple=True,
)
@click.option(
"--attach",
@ -1996,7 +2002,7 @@ def memory(
db.register_fts4_bm25()
if functions:
_register_functions(db, functions)
_maybe_register_functions(db, functions)
if return_db:
return db
@ -3281,6 +3287,13 @@ def _load_extensions(db, load_extension):
def _register_functions(db, functions):
# Register any Python functions as SQL functions:
# Check if this is a file path
if "\n" not in functions and functions.endswith(".py"):
try:
functions = pathlib.Path(functions).read_text()
except FileNotFoundError:
raise click.ClickException("File not found: {}".format(functions))
sqlite3.enable_callback_tracebacks(True)
globals = {}
try:
@ -3291,3 +3304,18 @@ def _register_functions(db, functions):
for name, value in globals.items():
if callable(value) and not name.startswith("_"):
db.register_function(value, name=name)
def _value_or_none(value):
if getattr(value, "__class__", None).__name__ == "Sentinel":
return None
return value
def _maybe_register_functions(db, functions_list):
if not functions_list:
return
for functions in functions_list:
functions = _value_or_none(functions)
if isinstance(functions, str) and functions.strip():
_register_functions(db, functions)