Also support def convert(value), closes #355

Plus added custom syntax error display

sqlite_utils/cli.py

@@ -17,6 +17,7 @@ import sys
 import csv as csv_std
 import tabulate
 from .utils import (
+    _compile_code,
     file_progress,
     find_spatialite,
     sqlite3,
@@ -2125,21 +2126,16 @@ def convert(
     if code == "-":
         # Read code from standard input
         code = sys.stdin.read()
-    # If single line and no 'return', add the return
-    if "\n" not in code and not code.strip().startswith("return "):
-        code = "return {}".format(code)
     where_args = dict(param) if param else []
     # Compile the code into a function body called fn(value)
-    new_code = ["def fn(value):"]
-    for line in code.split("\n"):
-        new_code.append("    {}".format(line))
-    code_o = compile("\n".join(new_code), "<string>", "exec")
-    locals = {}
-    globals = {"r": recipes, "recipes": recipes}
-    for import_ in imports:
-        globals[import_.split(".")[0]] = __import__(import_)
-    exec(code_o, globals, locals)
-    fn = locals["fn"]
+    try:
+        fn = _compile_code(code, imports)
+    except SyntaxError as e:
+        raise click.ClickException(
+            "Syntax error in code:\n\n{}\n\n{}".format(
+                textwrap.indent(e.text.strip(), "    "), e.msg
+            )
+        )
     if dry_run:
         # Pull first 20 values for first column and preview them
         db.conn.create_function("preview_transform", 1, lambda v: fn(v) if v else v)

sqlite_utils/utils.py

@@ -5,6 +5,7 @@ import enum
 import io
 import json
 import os
+from . import recipes
 from typing import cast, BinaryIO, Iterable, Optional, Tuple, Type
 
 import click
@@ -278,3 +279,30 @@ def progressbar(*args, **kwargs):
     else:
         with click.progressbar(*args, **kwargs) as bar:
             yield bar
+
+
+def _compile_code(code, imports):
+    locals = {}
+    globals = {"r": recipes, "recipes": recipes}
+    # If user defined a convert() function, return that
+    try:
+        exec(code, globals, locals)
+        return locals["convert"]
+    except (SyntaxError, NameError, KeyError, TypeError):
+        pass
+
+    # Try compiling their code as a function instead
+
+    # If single line and no 'return', add the return
+    if "\n" not in code and not code.strip().startswith("return "):
+        code = "return {}".format(code)
+
+    new_code = ["def fn(value):"]
+    for line in code.split("\n"):
+        new_code.append("    {}".format(line))
+    code_o = compile("\n".join(new_code), "<string>", "exec")
+
+    for import_ in imports:
+        globals[import_.split(".")[0]] = __import__(import_)
+    exec(code_o, globals, locals)
+    return locals["fn"]