From 22c8d10dd343476e8b7b9af3366fae4c8353dd2c Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Sun, 9 Jan 2022 12:06:02 -0800
Subject: [PATCH] --convert function can now modify row in place, closes #371

---
 docs/cli.rst              |  4 +---
 sqlite_utils/cli.py       |  8 ++------
 sqlite_utils/utils.py     | 24 +++++++++++++++++-------
 tests/test_cli_convert.py | 32 ++++++++++----------------------
 tests/test_cli_insert.py  | 19 +++++++++++++++++++
 5 files changed, 49 insertions(+), 38 deletions(-)

diff --git a/docs/cli.rst b/docs/cli.rst
index f56e382..233a4c0 100644
--- a/docs/cli.rst
+++ b/docs/cli.rst
@@ -959,9 +959,7 @@ Given a JSON file called ``dogs.json`` containing this:
 
 The following command will insert that data and add an ``is_good`` column set to ``1`` for each dog::
 
-    $ sqlite-utils insert dogs.db dogs dogs.json --convert '
-    row["is_good"] = 1
-    return row'
+    $ sqlite-utils insert dogs.db dogs dogs.json --convert 'row["is_good"] = 1`
 
 The ``--convert`` option also works with the ``--csv``, ``--tsv`` and ``--nl`` insert options.
 
diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py
index 2bd0278..212a039 100644
--- a/sqlite_utils/cli.py
+++ b/sqlite_utils/cli.py
@@ -824,7 +824,7 @@ def insert_upsert_implementation(
                 except TypeError:
                     raise click.ClickException("--convert must return dict or iterator")
         else:
-            docs = (fn(doc) for doc in docs)
+            docs = (fn(doc) or doc for doc in docs)
 
     extra_kwargs = {"ignore": ignore, "replace": replace, "truncate": truncate}
     if not_null:
@@ -2225,11 +2225,7 @@ def convert(
     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
-            )
-        )
+        raise click.ClickException(str(e))
     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)
diff --git a/sqlite_utils/utils.py b/sqlite_utils/utils.py
index 3b4769a..0777f30 100644
--- a/sqlite_utils/utils.py
+++ b/sqlite_utils/utils.py
@@ -293,15 +293,25 @@ def _compile_code(code, imports, variable="value"):
         pass
 
     # Try compiling their code as a function instead
-
-    # If single line and no 'return', add the return
+    body_variants = [code]
+    # If single line and no 'return', try adding the return
     if "\n" not in code and not code.strip().startswith("return "):
-        code = "return {}".format(code)
+        body_variants.insert(0, "return {}".format(code))
 
-    new_code = ["def fn({}):".format(variable)]
-    for line in code.split("\n"):
-        new_code.append("    {}".format(line))
-    code_o = compile("\n".join(new_code), "<string>", "exec")
+    code_o = None
+    for variant in body_variants:
+        new_code = ["def fn({}):".format(variable)]
+        for line in variant.split("\n"):
+            new_code.append("    {}".format(line))
+        try:
+            code_o = compile("\n".join(new_code), "<string>", "exec")
+            break
+        except SyntaxError:
+            # Try another variant, e.g. for 'return row["column"] = 1'
+            continue
+
+    if code_o is None:
+        raise SyntaxError("Could not compile code")
 
     for import_ in imports:
         globals[import_.split(".")[0]] = __import__(import_)
diff --git a/tests/test_cli_convert.py b/tests/test_cli_convert.py
index 481bebd..ec8110f 100644
--- a/tests/test_cli_convert.py
+++ b/tests/test_cli_convert.py
@@ -55,32 +55,20 @@ def test_convert_code(fresh_db_and_path, code):
 
 
 @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""",
-        ),
-    ],
+    "bad_code",
+    (
+        "def foo(value)",
+        "$",
+    ),
 )
-def test_convert_code_errors(fresh_db_and_path, bad_code, expected_error):
+def test_convert_code_errors(fresh_db_and_path, bad_code):
     db, db_path = fresh_db_and_path
     db["t"].insert({"text": "October"})
-    result = CliRunner().invoke(cli.cli, ["convert", db_path, "t", "text", bad_code])
+    result = CliRunner().invoke(
+        cli.cli, ["convert", db_path, "t", "text", bad_code], catch_exceptions=False
+    )
     assert 1 == result.exit_code
-    assert result.output.strip() == expected_error.strip()
+    assert result.output == "Error: Could not compile code\n"
 
 
 def test_convert_import(test_db_and_path):
diff --git a/tests/test_cli_insert.py b/tests/test_cli_insert.py
index 085f322..9c218cc 100644
--- a/tests/test_cli_insert.py
+++ b/tests/test_cli_insert.py
@@ -434,3 +434,22 @@ def test_insert_convert_lines(db_path):
     db = Database(db_path)
     rows = list(db.query("select [line] from [all]"))
     assert rows == [{"line": "THIS IS TEXT"}, {"line": "WILL BE UPPER NOW"}]
+
+
+def test_insert_convert_row_modifying_in_place(db_path):
+    result = CliRunner().invoke(
+        cli.cli,
+        [
+            "insert",
+            db_path,
+            "rows",
+            "-",
+            "--convert",
+            'row["is_chicken"] = True',
+        ],
+        input='{"name": "Azi"}',
+    )
+    assert result.exit_code == 0, result.output
+    db = Database(db_path)
+    rows = list(db.query("select name, is_chicken from rows"))
+    assert rows == [{"name": "Azi", "is_chicken": 1}]