Implemented --convert for different things, renamed --all to --text

This commit is contained in:
Simon Willison 2022-01-05 21:44:04 -08:00
commit 2e4847e493
4 changed files with 104 additions and 25 deletions

View file

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

View file

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

View file

@ -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), "<string>", "exec")

View file

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