--convert function can now modify row in place, closes #371

This commit is contained in:
Simon Willison 2022-01-09 12:06:02 -08:00
commit 22c8d10dd3
5 changed files with 49 additions and 38 deletions

View file

@ -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.

View file

@ -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)

View file

@ -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_)

View file

@ -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):

View file

@ -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}]