diff --git a/docs/cli.rst b/docs/cli.rst index 6a39c1e..aa11286 100644 --- a/docs/cli.rst +++ b/docs/cli.rst @@ -1029,17 +1029,24 @@ This supports nested imports as well, for example to use `ElementTree ", "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) diff --git a/sqlite_utils/utils.py b/sqlite_utils/utils.py index 00a3c02..fe71b73 100644 --- a/sqlite_utils/utils.py +++ b/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), "", "exec") + + for import_ in imports: + globals[import_.split(".")[0]] = __import__(import_) + exec(code_o, globals, locals) + return locals["fn"] diff --git a/tests/test_cli_convert.py b/tests/test_cli_convert.py index 8755600..481bebd 100644 --- a/tests/test_cli_convert.py +++ b/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):