From 2e4847e493a03d95f827ddfaa698c052e3b231a8 Mon Sep 17 00:00:00 2001 From: Simon Willison Date: Wed, 5 Jan 2022 21:44:04 -0800 Subject: [PATCH] Implemented --convert for different things, renamed --all to --text --- docs/cli.rst | 10 +++--- sqlite_utils/cli.py | 41 +++++++++++++++-------- sqlite_utils/utils.py | 6 ++-- tests/test_cli_insert.py | 72 +++++++++++++++++++++++++++++++++++++--- 4 files changed, 104 insertions(+), 25 deletions(-) diff --git a/docs/cli.rst b/docs/cli.rst index 5bf02d2..8ce2215 100644 --- a/docs/cli.rst +++ b/docs/cli.rst @@ -878,8 +878,8 @@ If you do this, the table will be created with column names called ``untitled_1` .. _cli_insert_unstructured: -Inserting unstructured data with \-\-lines and \-\-all -====================================================== +Inserting unstructured data with \-\-lines and \-\-text +======================================================= If you have an unstructured file you can insert its contents into a table with a single ``line`` column containing each line from the file using ``--lines``. This can be useful if you intend to further analyze those lines using SQL string functions or :ref:`sqlite-utils convert `:: @@ -893,16 +893,16 @@ This will produce the following schema: [line] TEXT ); -You can also insert the entire contents of the file into a single column called ``all`` using ``--all``:: +You can also insert the entire contents of the file into a single column called ``text`` using ``--text``:: - $ sqlite-utils insert content.db content file.txt --all + $ sqlite-utils insert content.db content file.txt --text The schema here will be: .. code-block:: sql CREATE TABLE [content] ( - [all] TEXT + [text] TEXT ); .. _cli_insert_replace: diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py index a5b979f..5b83dc5 100644 --- a/sqlite_utils/cli.py +++ b/sqlite_utils/cli.py @@ -657,7 +657,9 @@ def insert_upsert_options(fn): help="Treat each line as a single value called 'line'", ), click.option( - "--all", is_flag=True, help="Treat input as a single value called 'all'" + "--text", + is_flag=True, + help="Treat input as a single value called 'text'", ), click.option("--convert", help="Python code to convert each item"), click.option( @@ -723,7 +725,7 @@ def insert_upsert_implementation( csv, tsv, lines, - all, + text, convert, imports, delimiter, @@ -785,11 +787,8 @@ def insert_upsert_implementation( docs = tracker.wrap(docs) elif lines: docs = ({"line": line.strip()} for line in decoded) - elif all: - docs = ({"all": decoded.read()},) - elif convert: - fn = _compile_code(convert, imports) - docs = (fn(line) for line in decoded) + elif text: + docs = ({"text": decoded.read()},) else: try: if nl: @@ -805,6 +804,20 @@ def insert_upsert_implementation( if flatten: docs = (dict(_flatten(doc)) for doc in docs) + if convert: + variable = "row" + if lines: + variable = "line" + elif text: + variable = "text" + fn = _compile_code(convert, imports, variable=variable) + if lines: + docs = (fn(doc["line"]) for doc in docs) + elif text: + docs = (fn(doc["text"]) for doc in docs) + else: + docs = (fn(doc) for doc in docs) + extra_kwargs = {"ignore": ignore, "replace": replace, "truncate": truncate} if not_null: extra_kwargs["not_null"] = set(not_null) @@ -812,8 +825,10 @@ def insert_upsert_implementation( extra_kwargs["defaults"] = dict(default) if upsert: extra_kwargs["upsert"] = upsert + # Apply {"$base64": true, ...} decoding, if needed docs = (decode_base64_values(doc) for doc in docs) + try: db[table].insert_all( docs, pk=pk, batch_size=batch_size, alter=alter, **extra_kwargs @@ -889,7 +904,7 @@ def insert( csv, tsv, lines, - all, + text, convert, imports, delimiter, @@ -918,7 +933,7 @@ def insert( - Use --nl for newline-delimited JSON objects - Use --csv or --tsv for comma-separated or tab-separated input - Use --lines to write each incoming line to a column called "line" - - Use --all to write the entire input to a column called "all" + - Use --text to write the entire input to a column called "text" You can also use --convert to pass a fragment of Python code that will be used to convert each input. @@ -927,7 +942,7 @@ def insert( imported row, and can return a modified row. If you are using --lines your code will be passed a "line" variable, - and for --all an "all" variable. + and for --text an "text" variable. """ try: insert_upsert_implementation( @@ -940,7 +955,7 @@ def insert( csv, tsv, lines, - all, + text, convert, imports, delimiter, @@ -976,7 +991,7 @@ def upsert( csv, tsv, lines, - all, + text, convert, imports, batch_size, @@ -1008,7 +1023,7 @@ def upsert( csv, tsv, lines, - all, + text, convert, imports, delimiter, diff --git a/sqlite_utils/utils.py b/sqlite_utils/utils.py index fe71b73..caec976 100644 --- a/sqlite_utils/utils.py +++ b/sqlite_utils/utils.py @@ -281,14 +281,14 @@ def progressbar(*args, **kwargs): yield bar -def _compile_code(code, imports): +def _compile_code(code, imports, variable="value"): 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): + except (AttributeError, SyntaxError, NameError, KeyError, TypeError): pass # Try compiling their code as a function instead @@ -297,7 +297,7 @@ def _compile_code(code, imports): if "\n" not in code and not code.strip().startswith("return "): code = "return {}".format(code) - new_code = ["def fn(value):"] + new_code = ["def fn({}):".format(variable)] for line in code.split("\n"): new_code.append(" {}".format(line)) code_o = compile("\n".join(new_code), "", "exec") diff --git a/tests/test_cli_insert.py b/tests/test_cli_insert.py index 063c43f..bb9aa9b 100644 --- a/tests/test_cli_insert.py +++ b/tests/test_cli_insert.py @@ -339,14 +339,78 @@ def test_insert_lines(db_path): ] == list(db.query("select line from from_lines")) -def test_insert_all(db_path): +def test_insert_text(db_path): result = CliRunner().invoke( cli.cli, - ["insert", db_path, "from_all", "-", "--all"], + ["insert", db_path, "from_text", "-", "--text"], input='First line\nSecond line\n{"foo": "baz"}', ) assert 0 == result.exit_code, result.output db = Database(db_path) - assert [{"all": 'First line\nSecond line\n{"foo": "baz"}'}] == list( - db.query("select [all] from from_all") + assert [{"text": 'First line\nSecond line\n{"foo": "baz"}'}] == list( + db.query("select text from from_text") ) + + +@pytest.mark.parametrize( + "options,input", + ( + ([], '[{"id": "1", "name": "Bob"}, {"id": "2", "name": "Cat"}]'), + (["--csv"], "id,name\n1,Bob\n2,Cat"), + (["--nl"], '{"id": "1", "name": "Bob"}\n{"id": "2", "name": "Cat"}'), + ), +) +def test_insert_convert_json_csv_jsonnl(db_path, options, input): + result = CliRunner().invoke( + cli.cli, + ["insert", db_path, "rows", "-", "--convert", '{**row, **{"extra": 1}}'] + + options, + input=input, + ) + assert result.exit_code == 0, result.output + db = Database(db_path) + rows = list(db.query("select id, name, extra from rows")) + assert rows == [ + {"id": "1", "name": "Bob", "extra": 1}, + {"id": "2", "name": "Cat", "extra": 1}, + ] + + +def test_insert_convert_text(db_path): + result = CliRunner().invoke( + cli.cli, + [ + "insert", + db_path, + "text", + "-", + "--text", + "--convert", + '{"text": text.upper()}', + ], + input="This is text\nwill be upper now", + ) + assert result.exit_code == 0, result.output + db = Database(db_path) + rows = list(db.query("select [text] from [text]")) + assert rows == [{"text": "THIS IS TEXT\nWILL BE UPPER NOW"}] + + +def test_insert_convert_lines(db_path): + result = CliRunner().invoke( + cli.cli, + [ + "insert", + db_path, + "all", + "-", + "--lines", + "--convert", + '{"line": line.upper()}', + ], + input="This is text\nwill be upper now", + ) + assert result.exit_code == 0, result.output + db = Database(db_path) + rows = list(db.query("select [line] from [all]")) + assert rows == [{"line": "THIS IS TEXT"}, {"line": "WILL BE UPPER NOW"}]