Also support def convert(value), closes #355

Plus added custom syntax error display

docs/cli.rst

@@ -1029,17 +1029,24 @@ This supports nested imports as well, for example to use `ElementTree <https://d
         'xml.etree.ElementTree.fromstring(value).attrib["title"]' \
         --import=xml.etree.ElementTree
 
-Use a CODE value of `-` to read from standard input:
+Your code will be automatically wrapped in a function, but you can also define a function called `convert(value)` which will be called, if available::
+
+    $ sqlite-utils convert content.db articles headline '
+    def convert(value):
+        return value.upper()'
+
+Use a ``CODE`` value of ``-`` to read from standard input::
 
     $ cat mycode.py | sqlite-utils convert content.db articles headline -
 
-Where `mycode.py` contains a fragment of Python code that looks like this:
+Where ``mycode.py`` contains a fragment of Python code that looks like this:
 
-```python
-return value.upper()
-```
+.. code-block:: python
 
-The transformation will be applied to every row in the specified table. You can limit that to just rows that match a ``WHERE`` clause using ``--where``::
+    def convert(value):
+        return value.upper()
+
+The conversion will be applied to every row in the specified table. You can limit that to just rows that match a ``WHERE`` clause using ``--where``::
 
     $ sqlite-utils convert content.db articles headline 'value.upper()' \
         --where "headline like '%cat%'"

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"]

tests/test_cli_convert.py

@@ -35,39 +35,52 @@ def fresh_db_and_path(tmpdir):
         "return value.replace('October', 'Spooktober')",
         # Return is optional:
         "value.replace('October', 'Spooktober')",
+        # Multiple lines are supported:
+        "v = value.replace('October', 'Spooktober')\nreturn v",
+        # Can also define a convert() function
+        "def convert(value): return value.replace('October', 'Spooktober')",
+        # ... with imports
+        "import re\n\ndef convert(value): return value.replace('October', 'Spooktober')",
     ],
 )
-def test_convert_single_line(test_db_and_path, code):
-    db, db_path = test_db_and_path
-    result = CliRunner().invoke(cli.cli, ["convert", db_path, "example", "dt", code])
-    assert 0 == result.exit_code, result.output
-    assert [
-        {"id": 1, "dt": "5th Spooktober 2019 12:04"},
-        {"id": 2, "dt": "6th Spooktober 2019 00:05:06"},
-        {"id": 3, "dt": ""},
-        {"id": 4, "dt": None},
-    ] == list(db["example"].rows)
-
-
-def test_convert_multiple_lines(test_db_and_path):
-    db, db_path = test_db_and_path
+def test_convert_code(fresh_db_and_path, code):
+    db, db_path = fresh_db_and_path
+    db["t"].insert({"text": "October"})
     result = CliRunner().invoke(
-        cli.cli,
-        [
-            "convert",
-            db_path,
-            "example",
-            "dt",
-            "v = value.replace('October', 'Spooktober')\nreturn v.upper()",
-        ],
+        cli.cli, ["convert", db_path, "t", "text", code], catch_exceptions=False
     )
     assert 0 == result.exit_code, result.output
-    assert [
-        {"id": 1, "dt": "5TH SPOOKTOBER 2019 12:04"},
-        {"id": 2, "dt": "6TH SPOOKTOBER 2019 00:05:06"},
-        {"id": 3, "dt": ""},
-        {"id": 4, "dt": None},
-    ] == list(db["example"].rows)
+    value = list(db["t"].rows)[0]["text"]
+    assert value == "Spooktober"
+
+
+@pytest.mark.parametrize(
+    "bad_code,expected_error",
+    [
+        (
+            "def foo(value)",
+            """Error: Syntax error in code:
+
+    return def foo(value)
+
+invalid syntax""",
+        ),
+        (
+            "$",
+            """Error: Syntax error in code:
+
+    return $
+
+invalid syntax""",
+        ),
+    ],
+)
+def test_convert_code_errors(fresh_db_and_path, bad_code, expected_error):
+    db, db_path = fresh_db_and_path
+    db["t"].insert({"text": "October"})
+    result = CliRunner().invoke(cli.cli, ["convert", db_path, "t", "text", bad_code])
+    assert 1 == result.exit_code
+    assert result.output.strip() == expected_error.strip()
 
 
 def test_convert_import(test_db_and_path):