From 6964d67ce160e42202583d47177761e1a7093482 Mon Sep 17 00:00:00 2001 From: Simon Willison Date: Sat, 31 Jul 2021 21:33:00 -0700 Subject: [PATCH 01/18] Began implementation of sqlite-utils convert, refs #251 --- sqlite_utils/cli.py | 206 ++++++++++++++++++++++++++ sqlite_utils/utils.py | 14 ++ tests/test_convert.py | 331 ++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 551 insertions(+) create mode 100644 tests/test_convert.py diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py index b168ec0..56af19e 100644 --- a/sqlite_utils/cli.py +++ b/sqlite_utils/cli.py @@ -17,6 +17,7 @@ import tabulate from .utils import ( file_progress, find_spatialite, + progressbar, sqlite3, decode_base64_values, rows_from_file, @@ -1903,6 +1904,211 @@ def analyze_tables( click.echo(details) +@cli.command(name="convert") +@click.argument( + "db_path", + type=click.Path(file_okay=True, dir_okay=False, allow_dash=False), + required=True, +) +@click.argument("table", type=str) +@click.argument("columns", type=str, nargs=-1, required=True) +@click.argument("code", type=str) +@click.option( + "--import", "imports", type=str, multiple=True, help="Python modules to import" +) +@click.option( + "--dry-run", is_flag=True, help="Show results of running this against first 10 rows" +) +@click.option( + "--multi", is_flag=True, help="Populate columns for keys in returned dictionary" +) +@click.option("--output", help="Optional separate column to populate with the output") +@click.option( + "--output-type", + help="Column type to use for the output column", + default="text", + type=click.Choice(["integer", "float", "blob", "text"]), +) +@click.option("--drop", is_flag=True, help="Drop original column afterwards") +@click.option("-s", "--silent", is_flag=True, help="Don't show a progress bar") +def lambda_( + db_path, + table, + columns, + code, + imports, + dry_run, + multi, + output, + output_type, + drop, + silent, +): + """ + Convert columns using Python code you supply. For example: + + \b + $ sqlite-utils convert my.db mytable mycolumn + --code='"\\n".join(textwrap.wrap(value, 10))' + --import=textwrap + + "value" is a variable with the column value to be converted. + """ + if output is not None and len(columns) > 1: + raise click.ClickException("Cannot use --output with more than one column") + if multi and len(columns) > 1: + raise click.ClickException("Cannot use --multi with more than one column") + # If single line and no 'return', add the return + if "\n" not in code and not code.strip().startswith("return "): + code = "return {}".format(code) + # Compile the code into a function body called fn(value) + new_code = ["def fn(value):"] + for line in code.split("\n"): + new_code.append(" {}".format(line)) + code_o = compile("\n".join(new_code), "", "exec") + locals = {} + globals = {} + for import_ in imports: + globals[import_] = __import__(import_) + exec(code_o, globals, locals) + fn = locals["fn"] + if dry_run: + # Pull first 20 values for first column and preview them + db = sqlite3.connect(db_path) + db.create_function("preview_transform", 1, lambda v: fn(v) if v else v) + sql = """ + select + [{column}] as value, + preview_transform([{column}]) as preview + from [{table}] limit 10 + """.format( + column=columns[0], table=table + ) + for row in db.execute(sql).fetchall(): + print(row[0]) + print(" --- becomes:") + print(row[1]) + print() + elif multi: + _transform_multi(db_path, table, columns[0], fn, drop, silent) + else: + _transform(db_path, table, columns, fn, output, output_type, drop, silent) + + +def _transform(db_path, table, columns, fn, output, output_type, drop, silent): + db = sqlite_utils.Database(db_path) + count_sql = "select count(*) from [{}]".format(table) + todo_count = list(db.execute(count_sql).fetchall())[0][0] * len(columns) + + if drop and not output: + raise click.ClickException("--drop can only be used with --output or --multi") + + if output is not None: + if output not in db[table].columns_dict: + db[table].add_column(output, output_type or "text") + + with progressbar(length=todo_count, silent=silent) as bar: + + def transform_value(v): + bar.update(1) + if not v: + return v + return fn(v) + + db.register_function(transform_value) + sql = "update [{table}] set {sets};".format( + table=table, + sets=", ".join( + [ + "[{output_column}] = transform_value([{column}])".format( + output_column=output or column, column=column + ) + for column in columns + ] + ), + ) + with db.conn: + db.execute(sql) + if drop: + db[table].transform(drop=columns) + + +def _transform_multi(db_path, table, column, fn, drop, silent): + db = sqlite_utils.Database(db_path) + # First we execute the function + pk_to_values = {} + new_column_types = {} + pks = [column.name for column in db[table].columns if column.is_pk] + if not pks: + pks = ["rowid"] + + with progressbar( + length=db[table].count, silent=silent, label="1: Evaluating" + ) as bar: + for row in db[table].rows_where( + select=", ".join( + "[{}]".format(column_name) for column_name in (pks + [column]) + ) + ): + row_pk = tuple(row[pk] for pk in pks) + if len(row_pk) == 1: + row_pk = row_pk[0] + values = fn(row[column]) + if values is not None and not isinstance(values, dict): + raise click.ClickException( + "With --multi code must return a Python dictionary - returned {}".format( + repr(values) + ) + ) + if values: + for key, value in values.items(): + new_column_types.setdefault(key, set()).add(type(value)) + pk_to_values[row_pk] = values + bar.update(1) + + # Add any new columns + columns_to_create = _suggest_column_types(new_column_types) + for column_name, column_type in columns_to_create.items(): + if column_name not in db[table].columns_dict: + db[table].add_column(column_name, column_type) + + # Run the updates + with progressbar(length=db[table].count, silent=silent, label="2: Updating") as bar: + with db.conn: + for pk, updates in pk_to_values.items(): + db[table].update(pk, updates) + bar.update(1) + if drop: + db[table].transform(drop=(column,)) + + +def _suggest_column_types(all_column_types): + column_types = {} + for key, types in all_column_types.items(): + # Ignore null values if at least one other type present: + if len(types) > 1: + types.discard(None.__class__) + if {None.__class__} == types: + t = str + elif len(types) == 1: + t = list(types)[0] + # But if it's a subclass of list / tuple / dict, use str + # instead as we will be storing it as JSON in the table + for superclass in (list, tuple, dict): + if issubclass(t, superclass): + t = str + elif {int, bool}.issuperset(types): + t = int + elif {int, float, bool}.issuperset(types): + t = float + elif {bytes, str}.issuperset(types): + t = bytes + else: + t = str + column_types[key] = t + return column_types + + def _render_common(title, values): if values is None: return "" diff --git a/sqlite_utils/utils.py b/sqlite_utils/utils.py index 4f3c819..977dd79 100644 --- a/sqlite_utils/utils.py +++ b/sqlite_utils/utils.py @@ -254,3 +254,17 @@ class ValueTracker: not_these.append(name) for key in not_these: del self.couldbe[key] + + +class NullProgressBar: + def update(self, value): + pass + + +@contextlib.contextmanager +def progressbar(silent=False, **kwargs): + if silent: + yield NullProgressBar() + else: + with click.progressbar(**kwargs) as bar: + yield bar diff --git a/tests/test_convert.py b/tests/test_convert.py new file mode 100644 index 0000000..acaf555 --- /dev/null +++ b/tests/test_convert.py @@ -0,0 +1,331 @@ +from click.testing import CliRunner +from sqlite_utils import cli +import sqlite_utils +import textwrap +import pathlib +import pytest + + +@pytest.fixture +def test_db_and_path(fresh_db_and_path): + db, db_path = fresh_db_and_path + db["example"].insert_all( + [ + {"id": 1, "dt": "5th October 2019 12:04"}, + {"id": 2, "dt": "6th October 2019 00:05:06"}, + {"id": 3, "dt": ""}, + {"id": 4, "dt": None}, + ], + pk="id", + ) + return db, db_path + + +@pytest.fixture +def fresh_db_and_path(tmpdir): + db_path = str(pathlib.Path(tmpdir) / "data.db") + db = sqlite_utils.Database(db_path) + return db, db_path + + +@pytest.mark.parametrize( + "code", + [ + "return value.replace('October', 'Spooktober')", + # Return is optional: + "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 + result = CliRunner().invoke( + cli.cli, + [ + "convert", + db_path, + "example", + "dt", + "v = value.replace('October', 'Spooktober')\nreturn v.upper()", + ], + ) + 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_import(test_db_and_path): + db, db_path = test_db_and_path + result = CliRunner().invoke( + cli.cli, + [ + "convert", + db_path, + "example", + "dt", + "return re.sub('O..', 'OXX', value)", + "--import", + "re", + ], + ) + assert 0 == result.exit_code, result.output + assert [ + {"id": 1, "dt": "5th OXXober 2019 12:04"}, + {"id": 2, "dt": "6th OXXober 2019 00:05:06"}, + {"id": 3, "dt": ""}, + {"id": 4, "dt": None}, + ] == list(db["example"].rows) + + +def test_convert_dryrun(test_db_and_path): + db, db_path = test_db_and_path + result = CliRunner().invoke( + cli.cli, + [ + "convert", + db_path, + "example", + "dt", + "return re.sub('O..', 'OXX', value)", + "--import", + "re", + "--dry-run", + ], + ) + assert result.exit_code == 0 + assert result.output.strip() == ( + "5th October 2019 12:04\n" + " --- becomes:\n" + "5th OXXober 2019 12:04\n" + "\n" + "6th October 2019 00:05:06\n" + " --- becomes:\n" + "6th OXXober 2019 00:05:06\n" + "\n" + "\n" + " --- becomes:\n" + "\n" + "\n" + "None\n" + " --- becomes:\n" + "None" + ) + # But it should not have actually modified the table data + assert list(db["example"].rows) == [ + {"id": 1, "dt": "5th October 2019 12:04"}, + {"id": 2, "dt": "6th October 2019 00:05:06"}, + {"id": 3, "dt": ""}, + {"id": 4, "dt": None}, + ] + + +@pytest.mark.parametrize("drop", (True, False)) +def test_convert_output_column(test_db_and_path, drop): + db, db_path = test_db_and_path + args = [ + "convert", + db_path, + "example", + "dt", + "value.replace('October', 'Spooktober')", + "--output", + "newcol", + ] + if drop: + args += ["--drop"] + result = CliRunner().invoke(cli.cli, args) + assert 0 == result.exit_code, result.output + expected = [ + { + "id": 1, + "dt": "5th October 2019 12:04", + "newcol": "5th Spooktober 2019 12:04", + }, + { + "id": 2, + "dt": "6th October 2019 00:05:06", + "newcol": "6th Spooktober 2019 00:05:06", + }, + {"id": 3, "dt": "", "newcol": ""}, + {"id": 4, "dt": None, "newcol": None}, + ] + if drop: + for row in expected: + del row["dt"] + assert list(db["example"].rows) == expected + + +@pytest.mark.parametrize( + "output_type,expected", + ( + ("text", [(1, "1"), (2, "2"), (3, "3"), (4, "4")]), + ("float", [(1, 1.0), (2, 2.0), (3, 3.0), (4, 4.0)]), + ("integer", [(1, 1), (2, 2), (3, 3), (4, 4)]), + (None, [(1, "1"), (2, "2"), (3, "3"), (4, "4")]), + ), +) +def test_convert_output_column_output_type(test_db_and_path, output_type, expected): + db, db_path = test_db_and_path + args = [ + "convert", + db_path, + "example", + "id", + "value", + "--output", + "new_id", + ] + if output_type: + args += ["--output-type", output_type] + result = CliRunner().invoke( + cli.cli, + args, + ) + assert 0 == result.exit_code, result.output + assert expected == list(db.execute("select id, new_id from example")) + + +@pytest.mark.parametrize( + "options,expected_error", + [ + ( + [ + "dt", + "id", + "value.replace('October', 'Spooktober')", + "--output", + "newcol", + ], + "Cannot use --output with more than one column", + ), + ( + [ + "dt", + "value.replace('October', 'Spooktober')", + "--output", + "newcol", + "--output-type", + "invalid", + ], + "Error: Invalid value for '--output-type'", + ), + ( + [ + "value.replace('October', 'Spooktober')", + ], + "Missing argument 'COLUMNS...'", + ), + ], +) +def test_convert_output_error(test_db_and_path, options, expected_error): + db_path = test_db_and_path[1] + result = CliRunner().invoke( + cli.cli, + [ + "convert", + db_path, + "example", + ] + + options, + ) + assert result.exit_code != 0 + assert expected_error in result.output + + +@pytest.mark.parametrize("drop", (True, False)) +def test_convert_multi(fresh_db_and_path, drop): + db, db_path = fresh_db_and_path + db["creatures"].insert_all( + [ + {"id": 1, "name": "Simon"}, + {"id": 2, "name": "Cleo"}, + ], + pk="id", + ) + args = [ + "convert", + db_path, + "creatures", + "name", + "--multi", + '{"upper": value.upper(), "lower": value.lower()}', + ] + if drop: + args += ["--drop"] + result = CliRunner().invoke(cli.cli, args) + assert result.exit_code == 0, result.output + expected = [ + {"id": 1, "name": "Simon", "upper": "SIMON", "lower": "simon"}, + {"id": 2, "name": "Cleo", "upper": "CLEO", "lower": "cleo"}, + ] + if drop: + for row in expected: + del row["name"] + assert list(db["creatures"].rows) == expected + + +def test_convert_multi_complex_column_types(fresh_db_and_path): + db, db_path = fresh_db_and_path + db["rows"].insert_all( + [ + {"id": 1}, + {"id": 2}, + {"id": 3}, + {"id": 4}, + ], + pk="id", + ) + code = textwrap.dedent( + """ + if value == 1: + return {"is_str": "", "is_float": 1.2, "is_int": None} + elif value == 2: + return {"is_float": 1, "is_int": 12} + elif value == 3: + return {"is_bytes": b"blah"} + """ + ) + result = CliRunner().invoke( + cli.cli, + [ + "convert", + db_path, + "rows", + "id", + "--multi", + code, + ], + ) + assert result.exit_code == 0, result.output + assert list(db["rows"].rows) == [ + {"id": 1, "is_str": "", "is_float": 1.2, "is_int": None, "is_bytes": None}, + {"id": 2, "is_str": None, "is_float": 1.0, "is_int": 12, "is_bytes": None}, + { + "id": 3, + "is_str": None, + "is_float": None, + "is_int": None, + "is_bytes": b"blah", + }, + {"id": 4, "is_str": None, "is_float": None, "is_int": None, "is_bytes": None}, + ] + assert db["rows"].schema == ( + "CREATE TABLE [rows] (\n" + " [id] INTEGER PRIMARY KEY\n" + ", [is_str] TEXT, [is_float] FLOAT, [is_int] INTEGER, [is_bytes] BLOB)" + ) From 1b152951e58ce5aeacb30acc3faf8f60b2ab4516 Mon Sep 17 00:00:00 2001 From: Simon Willison Date: Sun, 1 Aug 2021 08:31:41 -0700 Subject: [PATCH 02/18] Initial documentation for sqlite-utils convert, refs #251 --- docs/cli.rst | 31 +++++++++++++++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/docs/cli.rst b/docs/cli.rst index d8993e4..d325900 100644 --- a/docs/cli.rst +++ b/docs/cli.rst @@ -920,6 +920,37 @@ The ``-`` argument indicates data should be read from standard input. The string When inserting data from standard input only the following column definitions are supported: ``name``, ``path``, ``content``, ``sha256``, ``md5`` and ``size``. +.. _cli_convert: + +Converting data in columns +========================== + +The ``convert`` command can be used to transform the data in a specified column - for example to parse a date string into an ISO timestamp, or to split a string of tags into a JSON array. + +The command accepts a database, table, one or more columns and a string of Python code to be executed against the values from those columns. The following example would replace the values in the ``headline`` column in the ``articles`` table with an upper-case version:: + + sqlite-utils convert content.db articles headline 'value.upper()' + +The Python code is passed as a string. Within that Python code the ``value`` variable will be the value of the current column. + +.. _cli_convert_output: + +Saving the result to a different column +--------------------------------------- + +The ``--output`` and ``--output-type`` options can be used to save the result of the conversion to a separate column, which will be created if that column does not already exist:: + + sqlite-utils convert content.db articles headline 'value.upper()' \ + --output headline_upper + +The type of the created column defaults to ``text``, but a different column type can be specified using ``--output-type``. This example will create a new floating point column called ``id_as_a_float`` with a copy of each item's ID increased by 0.5:: + + sqlite-utils convert content.db articles id 'float(value) + 0.5' \ + --output id_as_a_float \ + --output-type float + +You can drop the original column at the end of the operation by adding ``--drop``. + .. _cli_create_table: Creating tables From e5163a3b846d115739d0ffad671d7fb49b7df8bf Mon Sep 17 00:00:00 2001 From: Simon Willison Date: Sun, 1 Aug 2021 08:38:23 -0700 Subject: [PATCH 03/18] Fleshed out documentation for sqlite-utils convert a bit, refs #251 --- docs/cli.rst | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/docs/cli.rst b/docs/cli.rst index d325900..b6830f5 100644 --- a/docs/cli.rst +++ b/docs/cli.rst @@ -929,10 +929,24 @@ The ``convert`` command can be used to transform the data in a specified column The command accepts a database, table, one or more columns and a string of Python code to be executed against the values from those columns. The following example would replace the values in the ``headline`` column in the ``articles`` table with an upper-case version:: - sqlite-utils convert content.db articles headline 'value.upper()' + $ sqlite-utils convert content.db articles headline 'value.upper()' The Python code is passed as a string. Within that Python code the ``value`` variable will be the value of the current column. +The code you provide will be compiled into a function that takes ``value`` as a single argument. If you break your function body into multiple lines the last line should be a ``return`` statement:: + + $ sqlite-utils convert content.db articles headline ' + value = str(value) + return value.upper()' + +You can specify Python modules that should be imported and made available to your code using one or more ``--import`` options:: + + $ sqlite-utils convert content.db articles content \ + '"\n".join(textwrap.wrap(value, 10))' \ + --import=textwrap + +The ``--dry-run`` option will output a preview of the conversion against the first ten rows, without modifying the database. + .. _cli_convert_output: Saving the result to a different column From 58fd02a91cbf94bcdb33561274e4464d36d752b9 Mon Sep 17 00:00:00 2001 From: Simon Willison Date: Sun, 1 Aug 2021 08:47:26 -0700 Subject: [PATCH 04/18] Docs for sqlite-utils convert --multi, refs #251 --- docs/cli.rst | 34 ++++++++++++++++++++++++++++++++-- 1 file changed, 32 insertions(+), 2 deletions(-) diff --git a/docs/cli.rst b/docs/cli.rst index b6830f5..909c17e 100644 --- a/docs/cli.rst +++ b/docs/cli.rst @@ -954,17 +954,47 @@ Saving the result to a different column The ``--output`` and ``--output-type`` options can be used to save the result of the conversion to a separate column, which will be created if that column does not already exist:: - sqlite-utils convert content.db articles headline 'value.upper()' \ + $ sqlite-utils convert content.db articles headline 'value.upper()' \ --output headline_upper The type of the created column defaults to ``text``, but a different column type can be specified using ``--output-type``. This example will create a new floating point column called ``id_as_a_float`` with a copy of each item's ID increased by 0.5:: - sqlite-utils convert content.db articles id 'float(value) + 0.5' \ + $ sqlite-utils convert content.db articles id 'float(value) + 0.5' \ --output id_as_a_float \ --output-type float You can drop the original column at the end of the operation by adding ``--drop``. +.. _cli_convert_multi: + +Converting a column into multiple columns +----------------------------------------- + +Sometimes you may wish to convert a single column into multiple derived columns. For example, you may have a ``location`` column containing ``latitude,longitude`` values which you wish to split out into separate ``latitude`` and ``longitude`` columns. + +You can achieve this using the ``--multi`` option to ``sqlite-utils convert``. This option expects your Python code to return a Python dictionary: new columns well be created and populated for each of the keys in that dictionary. + +For the ``latitude,longitude`` example you would use the following:: + + $ sqlite-utils convert demo.db places location \ + 'bits = value.split(",") + return { + "latitude": float(bits[0]), + "longitude": float(bits[1]), + }' --multi + +The type of the returned values will be taken into account when creating the new columns. In this example, the resulting database schema will look like this: + +.. code-block:: sql + + CREATE TABLE [places] ( + [location] TEXT, + [latitude] FLOAT, + [longitude] FLOAT + ); + +The code function can also return ``None``, in which case its output will be ignored. You can drop the original column at the end of the operation by adding ``--drop``. + .. _cli_create_table: Creating tables From 89f754cd4be3b5f5dc100df0bdd982383fe10b23 Mon Sep 17 00:00:00 2001 From: Simon Willison Date: Sun, 1 Aug 2021 09:50:57 -0700 Subject: [PATCH 05/18] Updated sqlite-utils convert --help, refs #251 --- sqlite_utils/cli.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py index 56af19e..5a09a45 100644 --- a/sqlite_utils/cli.py +++ b/sqlite_utils/cli.py @@ -1931,7 +1931,7 @@ def analyze_tables( ) @click.option("--drop", is_flag=True, help="Drop original column afterwards") @click.option("-s", "--silent", is_flag=True, help="Don't show a progress bar") -def lambda_( +def convert( db_path, table, columns, @@ -1948,8 +1948,8 @@ def lambda_( Convert columns using Python code you supply. For example: \b - $ sqlite-utils convert my.db mytable mycolumn - --code='"\\n".join(textwrap.wrap(value, 10))' + $ sqlite-utils convert my.db mytable mycolumn \\ + '"\\n".join(textwrap.wrap(value, 10))' \\ --import=textwrap "value" is a variable with the column value to be converted. From 504ee7015a5e0d3f37ffce23977318043b8ea06a Mon Sep 17 00:00:00 2001 From: Simon Willison Date: Sun, 1 Aug 2021 12:58:47 -0700 Subject: [PATCH 06/18] Rename tests to test_cli_convert, refs #251 --- tests/{test_convert.py => test_cli_convert.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename tests/{test_convert.py => test_cli_convert.py} (100%) diff --git a/tests/test_convert.py b/tests/test_cli_convert.py similarity index 100% rename from tests/test_convert.py rename to tests/test_cli_convert.py From 53f9088963184c083e55d6402a6e2f643dc6b5c0 Mon Sep 17 00:00:00 2001 From: Simon Willison Date: Sun, 1 Aug 2021 14:01:41 -0700 Subject: [PATCH 07/18] Implemented recipes for sqlite-utils convert, refs #251 --- docs/cli.rst | 31 ++++++++++++++++ setup.py | 8 ++++- sqlite_utils/cli.py | 53 +++++++++++++++++++++------- sqlite_utils/recipes.py | 19 ++++++++++ tests/test_cli_convert.py | 74 +++++++++++++++++++++++++++++++++++++++ tests/test_docs.py | 33 +++++++++++++++-- 6 files changed, 202 insertions(+), 16 deletions(-) create mode 100644 sqlite_utils/recipes.py diff --git a/docs/cli.rst b/docs/cli.rst index 909c17e..ad752a6 100644 --- a/docs/cli.rst +++ b/docs/cli.rst @@ -947,6 +947,37 @@ You can specify Python modules that should be imported and made available to you The ``--dry-run`` option will output a preview of the conversion against the first ten rows, without modifying the database. +.. _cli_convert_recipes: + +sqlite-utils convert recipes +---------------------------- + +Various built-in recipe functions are available for common operations. These are: + +``r.jsonsplit(value, delimiter=',', type=)`` + Convert a string like ``a,b,c`` into a JSON array ``["a", "b", "c"]`` + + The ``delimiter`` parameter can be used to specify a different delimiter. + + The ``type`` parameter can be set to ``float`` or ``int`` to produce a JSON array of different types, for example if the column's string value was ``1.2,3,4`` the following:: + + r.jsonsplit(value, type=float) + + Would produce an array like this: ``[1.2, 3.0, 4.5]`` + +``r.parsedate(value, dayfirst=False, yearfirst=False)`` + Parse a date and convert it to ISO date format: ``yyyy-mm-dd`` + + In the case of dates such as ``03/04/05`` U.S. ``MM/DD/YY`` format is assumed - you can use ``dayfirst=True`` or ``yearfirst=True`` to change how these ambiguous dates are interpreted. + +``r.parsedatetime(value, dayfirst=False, yearfirst=False)`` + Parse a datetime and convert it to ISO datetime format: ``yyyy-mm-ddTHH:MM:SS`` + +These recipes can be used in the code passed to ``sqlite-utils convert`` like this:: + + $ sqlite-utils convert my.db mytable mycolumn \\ + 'r.jsonsplit(value, delimiter=":")' + .. _cli_convert_output: Saving the result to a different column diff --git a/setup.py b/setup.py index eeee83b..43e44ad 100644 --- a/setup.py +++ b/setup.py @@ -22,7 +22,13 @@ setup( version=VERSION, license="Apache License, Version 2.0", packages=find_packages(exclude=["tests", "tests.*"]), - install_requires=["sqlite-fts4", "click", "click-default-group", "tabulate"], + install_requires=[ + "sqlite-fts4", + "click", + "click-default-group", + "tabulate", + "dateutils", + ], setup_requires=["pytest-runner"], extras_require={ "test": ["pytest", "black", "hypothesis"], diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py index 5a09a45..d2e10cb 100644 --- a/sqlite_utils/cli.py +++ b/sqlite_utils/cli.py @@ -6,7 +6,9 @@ import hashlib import pathlib import sqlite_utils from sqlite_utils.db import AlterError, DescIndex +from sqlite_utils import recipes import textwrap +import inspect import io import itertools import json @@ -1904,7 +1906,43 @@ def analyze_tables( click.echo(details) -@cli.command(name="convert") +def _generate_convert_help(): + help = textwrap.dedent( + """ + Convert columns using Python code you supply. For example: + + \b + $ sqlite-utils convert my.db mytable mycolumn \\ + '"\\n".join(textwrap.wrap(value, 10))' \\ + --import=textwrap + + "value" is a variable with the column value to be converted. + + The following common operations are available as recipe functions: + """ + ).strip() + recipe_names = [ + n for n in dir(recipes) if not n.startswith("_") and n not in ("json", "parser") + ] + for name in recipe_names: + fn = getattr(recipes, name) + help += "\n\nr.{}{}\n\n {}".format( + name, str(inspect.signature(fn)), fn.__doc__ + ) + help += "\n\n" + help += textwrap.dedent( + """ + You can use these recipes like so: + + \b + $ sqlite-utils convert my.db mytable mycolumn \\ + 'r.jsonsplit(value, delimiter=":")' + """ + ).strip() + return help + + +@cli.command(help=_generate_convert_help()) @click.argument( "db_path", type=click.Path(file_okay=True, dir_okay=False, allow_dash=False), @@ -1944,16 +1982,7 @@ def convert( drop, silent, ): - """ - Convert columns using Python code you supply. For example: - - \b - $ sqlite-utils convert my.db mytable mycolumn \\ - '"\\n".join(textwrap.wrap(value, 10))' \\ - --import=textwrap - - "value" is a variable with the column value to be converted. - """ + sqlite3.enable_callback_tracebacks(True) if output is not None and len(columns) > 1: raise click.ClickException("Cannot use --output with more than one column") if multi and len(columns) > 1: @@ -1967,7 +1996,7 @@ def convert( new_code.append(" {}".format(line)) code_o = compile("\n".join(new_code), "", "exec") locals = {} - globals = {} + globals = {"r": recipes, "recipes": recipes} for import_ in imports: globals[import_] = __import__(import_) exec(code_o, globals, locals) diff --git a/sqlite_utils/recipes.py b/sqlite_utils/recipes.py new file mode 100644 index 0000000..6918661 --- /dev/null +++ b/sqlite_utils/recipes.py @@ -0,0 +1,19 @@ +from dateutil import parser +import json + + +def parsedate(value, dayfirst=False, yearfirst=False): + "Parse a date and convert it to ISO date format: yyyy-mm-dd" + return ( + parser.parse(value, dayfirst=dayfirst, yearfirst=yearfirst).date().isoformat() + ) + + +def parsedatetime(value, dayfirst=False, yearfirst=False): + "Parse a datetime and convert it to ISO datetime format: yyyy-mm-ddTHH:MM:SS" + return parser.parse(value, dayfirst=dayfirst, yearfirst=yearfirst).isoformat() + + +def jsonsplit(value, delimiter=",", type=str): + 'Convert a string like a,b,c into a JSON array ["a", "b", "c"]' + return json.dumps([type(s.strip()) for s in value.split(delimiter)]) diff --git a/tests/test_cli_convert.py b/tests/test_cli_convert.py index acaf555..1bb6c4f 100644 --- a/tests/test_cli_convert.py +++ b/tests/test_cli_convert.py @@ -1,6 +1,7 @@ from click.testing import CliRunner from sqlite_utils import cli import sqlite_utils +import json import textwrap import pathlib import pytest @@ -329,3 +330,76 @@ def test_convert_multi_complex_column_types(fresh_db_and_path): " [id] INTEGER PRIMARY KEY\n" ", [is_str] TEXT, [is_float] FLOAT, [is_int] INTEGER, [is_bytes] BLOB)" ) + + +@pytest.mark.parametrize("delimiter", [None, ";", "-"]) +def test_recipe_jsonsplit(tmpdir, delimiter): + db_path = str(pathlib.Path(tmpdir) / "data.db") + db = sqlite_utils.Database(db_path) + db["example"].insert_all( + [ + {"id": 1, "tags": (delimiter or ",").join(["foo", "bar"])}, + {"id": 2, "tags": (delimiter or ",").join(["bar", "baz"])}, + ], + pk="id", + ) + code = "r.jsonsplit(value)" + if delimiter: + code = 'recipes.jsonsplit(value, delimiter="{}")'.format(delimiter) + args = ["convert", db_path, "example", "tags", code] + result = CliRunner().invoke(cli.cli, args) + assert 0 == result.exit_code, result.output + assert list(db["example"].rows) == [ + {"id": 1, "tags": '["foo", "bar"]'}, + {"id": 2, "tags": '["bar", "baz"]'}, + ] + + +@pytest.mark.parametrize( + "type,expected_array", + ( + (None, ["1", "2", "3"]), + ("float", [1.0, 2.0, 3.0]), + ("int", [1, 2, 3]), + ), +) +def test_recipe_jsonsplit_type(fresh_db_and_path, type, expected_array): + db, db_path = fresh_db_and_path + db["example"].insert_all( + [ + {"id": 1, "records": "1,2,3"}, + ], + pk="id", + ) + code = "r.jsonsplit(value)" + if type: + code = "recipes.jsonsplit(value, type={})".format(type) + args = ["convert", db_path, "example", "records", code] + result = CliRunner().invoke(cli.cli, args) + assert 0 == result.exit_code, result.output + assert json.loads(db["example"].get(1)["records"]) == expected_array + + +@pytest.mark.parametrize("drop", (True, False)) +def test_recipe_jsonsplit_output(fresh_db_and_path, drop): + db, db_path = fresh_db_and_path + db["example"].insert_all( + [ + {"id": 1, "records": "1,2,3"}, + ], + pk="id", + ) + code = "r.jsonsplit(value)" + args = ["convert", db_path, "example", "records", code, "--output", "tags"] + if drop: + args += ["--drop"] + result = CliRunner().invoke(cli.cli, args) + assert 0 == result.exit_code, result.output + expected = { + "id": 1, + "records": "1,2,3", + "tags": '["1", "2", "3"]', + } + if drop: + del expected["records"] + assert db["example"].get(1) == expected diff --git a/tests/test_docs.py b/tests/test_docs.py index 87b685e..d760629 100644 --- a/tests/test_docs.py +++ b/tests/test_docs.py @@ -1,10 +1,12 @@ -from sqlite_utils import cli +from click.testing import CliRunner +from sqlite_utils import cli, recipes from pathlib import Path import pytest import re docs_path = Path(__file__).parent.parent / "docs" commands_re = re.compile(r"(?:\$ | )sqlite-utils (\S+) ") +recipes_re = re.compile(r"r\.(\w+)\(") @pytest.fixture(scope="session") @@ -17,11 +19,36 @@ def documented_commands(): } +@pytest.fixture(scope="session") +def documented_recipes(): + rst = (docs_path / "cli.rst").read_text() + return set(recipes_re.findall(rst)) + + @pytest.mark.parametrize("command", cli.cli.commands.keys()) def test_commands_are_documented(documented_commands, command): assert command in documented_commands @pytest.mark.parametrize("command", cli.cli.commands.values()) -def test_commands_have_docstrings(command): - assert command.__doc__, "{} is missing a docstring".format(command) +def test_commands_have_help(command): + assert command.help, "{} is missing its help".format(command) + + +def test_convert_help(): + result = CliRunner().invoke(cli.cli, ["convert", "--help"]) + assert result.exit_code == 0 + for expected in ( + "r.jsonsplit(value, ", + "r.parsedate(value, ", + "r.parsedatetime(value, ", + ): + assert expected in result.output + + +@pytest.mark.parametrize( + "recipe", + [n for n in dir(recipes) if not n.startswith("_") and n not in ("json", "parser")], +) +def test_recipes_are_documented(documented_recipes, recipe): + assert recipe in documented_recipes From dfed7ad18ae318d820e5fa87bd1931a4df5430a1 Mon Sep 17 00:00:00 2001 From: Simon Willison Date: Sun, 1 Aug 2021 14:09:59 -0700 Subject: [PATCH 08/18] Use types-python-dateutil --- setup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/setup.py b/setup.py index 43e44ad..5009736 100644 --- a/setup.py +++ b/setup.py @@ -33,7 +33,7 @@ setup( extras_require={ "test": ["pytest", "black", "hypothesis"], "docs": ["sphinx_rtd_theme", "sphinx-autobuild"], - "mypy": ["mypy", "types-click", "types-tabulate"], + "mypy": ["mypy", "types-click", "types-tabulate", "types-python-dateutil"], "flake8": ["flake8"], }, entry_points=""" From 5839ddb49017804bd724d3ab9dedc272ba2cf3dd Mon Sep 17 00:00:00 2001 From: Simon Willison Date: Sun, 1 Aug 2021 14:22:13 -0700 Subject: [PATCH 09/18] Do not allow --drop without --output or --multi --- sqlite_utils/cli.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py index d2e10cb..9b26e6a 100644 --- a/sqlite_utils/cli.py +++ b/sqlite_utils/cli.py @@ -1987,6 +1987,8 @@ def convert( raise click.ClickException("Cannot use --output with more than one column") if multi and len(columns) > 1: raise click.ClickException("Cannot use --multi with more than one column") + if drop and not (output or multi): + raise click.ClickException("--drop can only be used with --output or --multi") # If single line and no 'return', add the return if "\n" not in code and not code.strip().startswith("return "): code = "return {}".format(code) From eef46da8f54d3e47490bd9322fe796f0a1e6bbe3 Mon Sep 17 00:00:00 2001 From: Simon Willison Date: Sun, 1 Aug 2021 14:22:41 -0700 Subject: [PATCH 10/18] table.convert(...) method, refs #251 --- sqlite_utils/db.py | 50 +++++++++++++++++++++++++++++++++++++++++- tests/test_convert.py | 51 +++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 100 insertions(+), 1 deletion(-) create mode 100644 tests/test_convert.py diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index ae87e27..03cea3f 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -1,4 +1,10 @@ -from .utils import sqlite3, OperationalError, suggest_column_types, column_affinity +from .utils import ( + sqlite3, + OperationalError, + suggest_column_types, + column_affinity, + progressbar, +) from collections import namedtuple from collections.abc import Mapping import contextlib @@ -1697,6 +1703,48 @@ class Table(Queryable): self.last_pk = pk_values[0] if len(pks) == 1 else pk_values return self + def convert( + self, + columns, + fn, + output=None, + output_type=None, + show_progress=False, + drop=False, + ): + if isinstance(columns, str): + columns = [columns] + todo_count = self.count * len(columns) + + if output is not None: + if output not in self.columns_dict: + self.add_column(output, output_type or "text") + + with progressbar(length=todo_count, silent=not show_progress) as bar: + + def transform_value(v): + bar.update(1) + if not v: + return v + return fn(v) + + self.db.register_function(transform_value) + sql = "update [{table}] set {sets};".format( + table=self.name, + sets=", ".join( + [ + "[{output_column}] = transform_value([{column}])".format( + output_column=output or column, column=column + ) + for column in columns + ] + ), + ) + with self.db.conn: + self.db.execute(sql) + if drop: + self.transform(drop=columns) + def build_insert_queries_and_params( self, extracts, diff --git a/tests/test_convert.py b/tests/test_convert.py new file mode 100644 index 0000000..739195d --- /dev/null +++ b/tests/test_convert.py @@ -0,0 +1,51 @@ +import pytest + + +@pytest.mark.parametrize( + "columns,fn,expected", + ( + ( + "title", + lambda value: value.upper(), + {"title": "MIXED CASE", "abstract": "Abstract"}, + ), + ( + ["title", "abstract"], + lambda value: value.upper(), + {"title": "MIXED CASE", "abstract": "ABSTRACT"}, + ), + ), +) +def test_convert(fresh_db, columns, fn, expected): + table = fresh_db["table"] + table.insert({"title": "Mixed Case", "abstract": "Abstract"}) + table.convert(columns, fn) + assert list(table.rows) == [expected] + + +@pytest.mark.parametrize( + "drop,expected", + ( + (False, {"title": "Mixed Case", "other": "MIXED CASE"}), + (True, {"other": "MIXED CASE"}), + ), +) +def test_convert_output(fresh_db, drop, expected): + table = fresh_db["table"] + table.insert({"title": "Mixed Case"}) + table.convert("title", lambda v: v.upper(), output="other", drop=drop) + assert list(table.rows) == [expected] + + +@pytest.mark.parametrize( + "type,expected", + ( + (int, {"other": 123}), + (float, {"other": 123.0}), + ), +) +def test_convert_output_type(fresh_db, type, expected): + table = fresh_db["table"] + table.insert({"number": "123"}) + table.convert("number", lambda v: v, output="other", output_type=type, drop=True) + assert list(table.rows) == [expected] From 3532bcca9776fd4f7051d8dbd59c74be9574620a Mon Sep 17 00:00:00 2001 From: Simon Willison Date: Sun, 1 Aug 2021 14:31:06 -0700 Subject: [PATCH 11/18] Test for --drop error without --output or --multi --- tests/test_cli_convert.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/tests/test_cli_convert.py b/tests/test_cli_convert.py index 1bb6c4f..d3f5d28 100644 --- a/tests/test_cli_convert.py +++ b/tests/test_cli_convert.py @@ -403,3 +403,10 @@ def test_recipe_jsonsplit_output(fresh_db_and_path, drop): if drop: del expected["records"] assert db["example"].get(1) == expected + + +def test_cannot_use_drop_without_multi_or_output(fresh_db_and_path): + args = ["convert", fresh_db_and_path[1], "example", "records", "value", "--drop"] + result = CliRunner().invoke(cli.cli, args) + assert result.exit_code == 1, result.output + assert "Error: --drop can only be used with --output or --multi" in result.output From d0e25fbc0d802bb5724730f21c92de04a6f1d794 Mon Sep 17 00:00:00 2001 From: Simon Willison Date: Sun, 1 Aug 2021 14:40:35 -0700 Subject: [PATCH 12/18] Tests for recipes --- tests/test_recipes.py | 103 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 103 insertions(+) create mode 100644 tests/test_recipes.py diff --git a/tests/test_recipes.py b/tests/test_recipes.py new file mode 100644 index 0000000..c596487 --- /dev/null +++ b/tests/test_recipes.py @@ -0,0 +1,103 @@ +from sqlite_utils import recipes +from sqlite_utils.utils import sqlite3 +import json +import pytest + + +@pytest.fixture +def dates_db(fresh_db): + fresh_db["example"].insert_all( + [ + {"id": 1, "dt": "5th October 2019 12:04"}, + {"id": 2, "dt": "6th October 2019 00:05:06"}, + {"id": 3, "dt": ""}, + {"id": 4, "dt": None}, + ], + pk="id", + ) + return fresh_db + + +def test_parsedate(dates_db): + dates_db["example"].convert("dt", recipes.parsedate) + assert list(dates_db["example"].rows) == [ + {"id": 1, "dt": "2019-10-05"}, + {"id": 2, "dt": "2019-10-06"}, + {"id": 3, "dt": ""}, + {"id": 4, "dt": None}, + ] + + +def test_parsedatetime(dates_db): + dates_db["example"].convert("dt", recipes.parsedatetime) + assert list(dates_db["example"].rows) == [ + {"id": 1, "dt": "2019-10-05T12:04:00"}, + {"id": 2, "dt": "2019-10-06T00:05:06"}, + {"id": 3, "dt": ""}, + {"id": 4, "dt": None}, + ] + + +@pytest.mark.parametrize( + "recipe,kwargs,expected", + ( + ("parsedate", {}, "2005-03-04"), + ("parsedate", {"dayfirst": True}, "2005-04-03"), + ("parsedatetime", {}, "2005-03-04T00:00:00"), + ("parsedatetime", {"dayfirst": True}, "2005-04-03T00:00:00"), + ), +) +def test_dayfirst_yearfirst(fresh_db, recipe, kwargs, expected): + fresh_db["example"].insert_all( + [ + {"id": 1, "dt": "03/04/05"}, + ], + pk="id", + ) + fresh_db["example"].convert( + "dt", lambda value: getattr(recipes, recipe)(value, **kwargs) + ) + assert list(fresh_db["example"].rows) == [ + {"id": 1, "dt": expected}, + ] + + +@pytest.mark.parametrize("delimiter", [None, ";", "-"]) +def test_jsonsplit(fresh_db, delimiter): + fresh_db["example"].insert_all( + [ + {"id": 1, "tags": (delimiter or ",").join(["foo", "bar"])}, + {"id": 2, "tags": (delimiter or ",").join(["bar", "baz"])}, + ], + pk="id", + ) + fn = recipes.jsonsplit + if delimiter is not None: + fn = lambda value: recipes.jsonsplit(value, delimiter=delimiter) + fresh_db["example"].convert("tags", fn) + assert list(fresh_db["example"].rows) == [ + {"id": 1, "tags": '["foo", "bar"]'}, + {"id": 2, "tags": '["bar", "baz"]'}, + ] + + +@pytest.mark.parametrize( + "type,expected", + ( + (None, ["1", "2", "3"]), + (float, [1.0, 2.0, 3.0]), + (int, [1, 2, 3]), + ), +) +def test_jsonsplit_type(fresh_db, type, expected): + fresh_db["example"].insert_all( + [ + {"id": 1, "records": "1,2,3"}, + ], + pk="id", + ) + fn = recipes.jsonsplit + if type is not None: + fn = lambda value: recipes.jsonsplit(value, type=type) + fresh_db["example"].convert("records", fn) + assert json.loads(fresh_db["example"].get(1)["records"]) == expected From 570bee7edd17424e37ce7d7e6a1b35d4be063f17 Mon Sep 17 00:00:00 2001 From: Simon Willison Date: Sun, 1 Aug 2021 21:08:40 -0700 Subject: [PATCH 13/18] Implemented .convert(multi=True) and refactored CLI to use that method, refs #251, #302 --- sqlite_utils/cli.py | 146 +++++++----------------------------------- sqlite_utils/db.py | 64 +++++++++++++++++- sqlite_utils/utils.py | 5 +- 3 files changed, 90 insertions(+), 125 deletions(-) diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py index 9b26e6a..319dad6 100644 --- a/sqlite_utils/cli.py +++ b/sqlite_utils/cli.py @@ -5,7 +5,7 @@ from datetime import datetime import hashlib import pathlib import sqlite_utils -from sqlite_utils.db import AlterError, DescIndex +from sqlite_utils.db import AlterError, BadMultiValues, DescIndex from sqlite_utils import recipes import textwrap import inspect @@ -23,6 +23,7 @@ from .utils import ( sqlite3, decode_base64_values, rows_from_file, + types_for_column_types, Format, TypeTracker, ) @@ -1983,6 +1984,7 @@ def convert( silent, ): sqlite3.enable_callback_tracebacks(True) + db = sqlite_utils.Database(db_path) if output is not None and len(columns) > 1: raise click.ClickException("Cannot use --output with more than one column") if multi and len(columns) > 1: @@ -2005,8 +2007,7 @@ def convert( fn = locals["fn"] if dry_run: # Pull first 20 values for first column and preview them - db = sqlite3.connect(db_path) - db.create_function("preview_transform", 1, lambda v: fn(v) if v else v) + db.conn.create_function("preview_transform", 1, lambda v: fn(v) if v else v) sql = """ select [{column}] as value, @@ -2015,129 +2016,28 @@ def convert( """.format( column=columns[0], table=table ) - for row in db.execute(sql).fetchall(): - print(row[0]) - print(" --- becomes:") - print(row[1]) - print() - elif multi: - _transform_multi(db_path, table, columns[0], fn, drop, silent) + for row in db.conn.execute(sql).fetchall(): + click.echo(str(row[0])) + click.echo(" --- becomes:") + click.echo(str(row[1])) + click.echo() else: - _transform(db_path, table, columns, fn, output, output_type, drop, silent) - - -def _transform(db_path, table, columns, fn, output, output_type, drop, silent): - db = sqlite_utils.Database(db_path) - count_sql = "select count(*) from [{}]".format(table) - todo_count = list(db.execute(count_sql).fetchall())[0][0] * len(columns) - - if drop and not output: - raise click.ClickException("--drop can only be used with --output or --multi") - - if output is not None: - if output not in db[table].columns_dict: - db[table].add_column(output, output_type or "text") - - with progressbar(length=todo_count, silent=silent) as bar: - - def transform_value(v): - bar.update(1) - if not v: - return v - return fn(v) - - db.register_function(transform_value) - sql = "update [{table}] set {sets};".format( - table=table, - sets=", ".join( - [ - "[{output_column}] = transform_value([{column}])".format( - output_column=output or column, column=column - ) - for column in columns - ] - ), - ) - with db.conn: - db.execute(sql) - if drop: - db[table].transform(drop=columns) - - -def _transform_multi(db_path, table, column, fn, drop, silent): - db = sqlite_utils.Database(db_path) - # First we execute the function - pk_to_values = {} - new_column_types = {} - pks = [column.name for column in db[table].columns if column.is_pk] - if not pks: - pks = ["rowid"] - - with progressbar( - length=db[table].count, silent=silent, label="1: Evaluating" - ) as bar: - for row in db[table].rows_where( - select=", ".join( - "[{}]".format(column_name) for column_name in (pks + [column]) + try: + db[table].convert( + columns, + fn, + output=output, + output_type=output_type, + drop=drop, + multi=multi, + show_progress=not silent, ) - ): - row_pk = tuple(row[pk] for pk in pks) - if len(row_pk) == 1: - row_pk = row_pk[0] - values = fn(row[column]) - if values is not None and not isinstance(values, dict): - raise click.ClickException( - "With --multi code must return a Python dictionary - returned {}".format( - repr(values) - ) + except BadMultiValues as e: + raise click.ClickException( + "When using --multi code must return a Python dictionary - returned: {}".format( + repr(e.values) ) - if values: - for key, value in values.items(): - new_column_types.setdefault(key, set()).add(type(value)) - pk_to_values[row_pk] = values - bar.update(1) - - # Add any new columns - columns_to_create = _suggest_column_types(new_column_types) - for column_name, column_type in columns_to_create.items(): - if column_name not in db[table].columns_dict: - db[table].add_column(column_name, column_type) - - # Run the updates - with progressbar(length=db[table].count, silent=silent, label="2: Updating") as bar: - with db.conn: - for pk, updates in pk_to_values.items(): - db[table].update(pk, updates) - bar.update(1) - if drop: - db[table].transform(drop=(column,)) - - -def _suggest_column_types(all_column_types): - column_types = {} - for key, types in all_column_types.items(): - # Ignore null values if at least one other type present: - if len(types) > 1: - types.discard(None.__class__) - if {None.__class__} == types: - t = str - elif len(types) == 1: - t = list(types)[0] - # But if it's a subclass of list / tuple / dict, use str - # instead as we will be storing it as JSON in the table - for superclass in (list, tuple, dict): - if issubclass(t, superclass): - t = str - elif {int, bool}.issuperset(types): - t = int - elif {int, float, bool}.issuperset(types): - t = float - elif {bytes, str}.issuperset(types): - t = bytes - else: - t = str - column_types[key] = t - return column_types + ) def _render_common(title, values): diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index 03cea3f..5f9a526 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -2,11 +2,13 @@ from .utils import ( sqlite3, OperationalError, suggest_column_types, + types_for_column_types, column_affinity, progressbar, ) from collections import namedtuple from collections.abc import Mapping +import click import contextlib import datetime import decimal @@ -159,6 +161,13 @@ class DescIndex(str): pass +class BadMultiValues(Exception): + "With multi=True code must return a Python dictionary" + + def __init__(self, values): + self.values = values + + _COUNTS_TABLE_CREATE_SQL = """ CREATE TABLE IF NOT EXISTS [{}]( [table] TEXT PRIMARY KEY, @@ -1709,11 +1718,18 @@ class Table(Queryable): fn, output=None, output_type=None, - show_progress=False, drop=False, + multi=False, + show_progress=False, ): if isinstance(columns, str): columns = [columns] + + if multi: + return self._convert_multi( + columns[0], fn, drop=drop, show_progress=show_progress + ) + todo_count = self.count * len(columns) if output is not None: @@ -1744,6 +1760,52 @@ class Table(Queryable): self.db.execute(sql) if drop: self.transform(drop=columns) + return self + + def _convert_multi(self, column, fn, drop, show_progress): + # First we execute the function + pk_to_values = {} + new_column_types = {} + pks = [column.name for column in self.columns if column.is_pk] + if not pks: + pks = ["rowid"] + + with progressbar( + length=self.count, silent=not show_progress, label="1: Evaluating" + ) as bar: + for row in self.rows_where( + select=", ".join( + "[{}]".format(column_name) for column_name in (pks + [column]) + ) + ): + row_pk = tuple(row[pk] for pk in pks) + if len(row_pk) == 1: + row_pk = row_pk[0] + values = fn(row[column]) + if values is not None and not isinstance(values, dict): + raise BadMultiValues(values) + if values: + for key, value in values.items(): + new_column_types.setdefault(key, set()).add(type(value)) + pk_to_values[row_pk] = values + bar.update(1) + + # Add any new columns + columns_to_create = types_for_column_types(new_column_types) + for column_name, column_type in columns_to_create.items(): + if column_name not in self.columns_dict: + self.add_column(column_name, column_type) + + # Run the updates + with progressbar( + length=self.count, silent=not show_progress, label="2: Updating" + ) as bar: + with self.db.conn: + for pk, updates in pk_to_values.items(): + self.update(pk, updates) + bar.update(1) + if drop: + self.transform(drop=(column,)) def build_insert_queries_and_params( self, diff --git a/sqlite_utils/utils.py b/sqlite_utils/utils.py index 977dd79..a781469 100644 --- a/sqlite_utils/utils.py +++ b/sqlite_utils/utils.py @@ -31,8 +31,11 @@ def suggest_column_types(records): for record in records: for key, value in record.items(): all_column_types.setdefault(key, set()).add(type(value)) - column_types = {} + return types_for_column_types(all_column_types) + +def types_for_column_types(all_column_types): + column_types = {} for key, types in all_column_types.items(): # Ignore null values if at least one other type present: if len(types) > 1: From d2436c148eb3d92f56b41666dc9b7948d2d8f389 Mon Sep 17 00:00:00 2001 From: Simon Willison Date: Sun, 1 Aug 2021 21:11:40 -0700 Subject: [PATCH 14/18] Unit tests for .convert(multi=True), refs #251, #302 --- tests/test_convert.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/tests/test_convert.py b/tests/test_convert.py index 739195d..15e8026 100644 --- a/tests/test_convert.py +++ b/tests/test_convert.py @@ -1,3 +1,4 @@ +from sqlite_utils.db import BadMultiValues import pytest @@ -49,3 +50,21 @@ def test_convert_output_type(fresh_db, type, expected): table.insert({"number": "123"}) table.convert("number", lambda v: v, output="other", output_type=type, drop=True) assert list(table.rows) == [expected] + + +def test_convert_multi(fresh_db): + table = fresh_db["table"] + table.insert({"title": "Mixed Case"}) + table.convert( + "title", lambda v: {"upper": v.upper(), "lower": v.lower()}, multi=True + ) + assert list(table.rows) == [ + {"title": "Mixed Case", "upper": "MIXED CASE", "lower": "mixed case"} + ] + + +def test_convert_multi_exception(fresh_db): + table = fresh_db["table"] + table.insert({"title": "Mixed Case"}) + with pytest.raises(BadMultiValues): + table.convert("title", lambda v: v.upper(), multi=True) From da93cf115e7a8f21d123dd65d884e818bc42e126 Mon Sep 17 00:00:00 2001 From: Simon Willison Date: Sun, 1 Aug 2021 21:17:53 -0700 Subject: [PATCH 15/18] Test more sqlite-utils convert error conditions --- tests/test_cli_convert.py | 29 +++++++++++++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/tests/test_cli_convert.py b/tests/test_cli_convert.py index d3f5d28..51634cd 100644 --- a/tests/test_cli_convert.py +++ b/tests/test_cli_convert.py @@ -410,3 +410,32 @@ def test_cannot_use_drop_without_multi_or_output(fresh_db_and_path): result = CliRunner().invoke(cli.cli, args) assert result.exit_code == 1, result.output assert "Error: --drop can only be used with --output or --multi" in result.output + + +def test_cannot_use_multi_with_more_than_one_column(fresh_db_and_path): + args = [ + "convert", + fresh_db_and_path[1], + "example", + "records", + "othercol", + "value", + "--multi", + ] + result = CliRunner().invoke(cli.cli, args) + assert result.exit_code == 1, result.output + assert "Error: Cannot use --multi with more than one column" in result.output + + +def test_multi_with_bad_function(test_db_and_path): + args = [ + "convert", + test_db_and_path[1], + "example", + "dt", + "value.upper()", + "--multi", + ] + result = CliRunner().invoke(cli.cli, args) + assert result.exit_code == 1, result.output + assert "When using --multi code must return a Python dictionary" in result.output From 93114feefb943adf8616e04b690a5a82050d3944 Mon Sep 17 00:00:00 2001 From: Simon Willison Date: Sun, 1 Aug 2021 21:21:21 -0700 Subject: [PATCH 16/18] flake8 fixes --- sqlite_utils/cli.py | 2 -- sqlite_utils/db.py | 1 - tests/test_recipes.py | 11 ++++++++--- 3 files changed, 8 insertions(+), 6 deletions(-) diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py index 319dad6..e1b770f 100644 --- a/sqlite_utils/cli.py +++ b/sqlite_utils/cli.py @@ -19,11 +19,9 @@ import tabulate from .utils import ( file_progress, find_spatialite, - progressbar, sqlite3, decode_base64_values, rows_from_file, - types_for_column_types, Format, TypeTracker, ) diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index 5f9a526..59c6998 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -8,7 +8,6 @@ from .utils import ( ) from collections import namedtuple from collections.abc import Mapping -import click import contextlib import datetime import decimal diff --git a/tests/test_recipes.py b/tests/test_recipes.py index c596487..89240a2 100644 --- a/tests/test_recipes.py +++ b/tests/test_recipes.py @@ -1,5 +1,4 @@ from sqlite_utils import recipes -from sqlite_utils.utils import sqlite3 import json import pytest @@ -73,7 +72,10 @@ def test_jsonsplit(fresh_db, delimiter): ) fn = recipes.jsonsplit if delimiter is not None: - fn = lambda value: recipes.jsonsplit(value, delimiter=delimiter) + + def fn(value): + return recipes.jsonsplit(value, delimiter=delimiter) + fresh_db["example"].convert("tags", fn) assert list(fresh_db["example"].rows) == [ {"id": 1, "tags": '["foo", "bar"]'}, @@ -98,6 +100,9 @@ def test_jsonsplit_type(fresh_db, type, expected): ) fn = recipes.jsonsplit if type is not None: - fn = lambda value: recipes.jsonsplit(value, type=type) + + def fn(value): + return recipes.jsonsplit(value, type=type) + fresh_db["example"].convert("records", fn) assert json.loads(fresh_db["example"].get(1)["records"]) == expected From 4c3bf9730542f1e49cbf11a61dbdb7fd621453df Mon Sep 17 00:00:00 2001 From: Simon Willison Date: Sun, 1 Aug 2021 21:43:20 -0700 Subject: [PATCH 17/18] Documentation for table.convert(...), refs #302 --- docs/cli.rst | 2 +- docs/python-api.rst | 41 ++++++++++++++++++++++++++++++++++++++++- sqlite_utils/db.py | 10 +++++----- tests/test_convert.py | 7 +++++++ 4 files changed, 53 insertions(+), 7 deletions(-) diff --git a/docs/cli.rst b/docs/cli.rst index ad752a6..7bc7a3b 100644 --- a/docs/cli.rst +++ b/docs/cli.rst @@ -975,7 +975,7 @@ Various built-in recipe functions are available for common operations. These are These recipes can be used in the code passed to ``sqlite-utils convert`` like this:: - $ sqlite-utils convert my.db mytable mycolumn \\ + $ sqlite-utils convert my.db mytable mycolumn \ 'r.jsonsplit(value, delimiter=":")' .. _cli_convert_output: diff --git a/docs/python-api.rst b/docs/python-api.rst index 7f3cdd5..7f490ab 100644 --- a/docs/python-api.rst +++ b/docs/python-api.rst @@ -702,7 +702,7 @@ You can delete all records in a table that match a specific WHERE statement usin >>> db = sqlite_utils.Database("dogs.db") >>> # Delete every dog with age less than 3 - >>> db["dogs"].delete_where("age < ?", [3]): + >>> db["dogs"].delete_where("age < ?", [3]) Calling ``table.delete_where()`` with no other arguments will delete every row in the table. @@ -736,6 +736,45 @@ An ``upsert_all()`` method is also available, which behaves like ``insert_all()` .. note:: ``.upsert()`` and ``.upsert_all()`` in sqlite-utils 1.x worked like ``.insert(..., replace=True)`` and ``.insert_all(..., replace=True)`` do in 2.x. See `issue #66 `__ for details of this change. +.. _python_api_convert: + +Converting data in columns +========================== + +The ``table.convert(...)`` method can be used to apply a conversion function to the values in a column, either to update that column or to populate new columns. It is the Python library equivalent of the :ref:`sqlite-utils convert ` command. + +This feature works by registering a custom SQLite function that applies a Python transformation, then running a SQL query equivalent to ``UPDATE table SET column = convert_value(column);`` + +To transform a specific column to uppercase, you would use the following: + +.. code-block:: python + + db["dogs"].convert("name", lambda value: value.upper()) + +You can pass a list of columns, in which case the transformation will be applied to each one: + +.. code-block:: python + + db["dogs"].convert(["name", "twitter"], lambda value: value.upper()) + +To save the output to of the transformation to a different column, use the ``output=`` parameter: + +.. code-block:: python + + db["dogs"].convert("name", lambda value: value.upper(), output="name_upper") + +This will add the new column, if it does not already exist. You can pass ``output_type=int`` or some other type to control the type of the new column - otherwise it will default to text. + +If you want to drop the original column after saving the results in a separate output column, pass ``drop=True``. + +You can create multiple new columns from a single input column by passing ``multi=True`` and a conversion function that returns a Python dictionary. This example creates new ``upper`` and ``lower`` columns populated from the single ``title`` column: + +.. code-block:: python + + table.convert( + "title", lambda v: {"upper": v.upper(), "lower": v.lower()}, multi=True + ) + .. _python_api_lookup_tables: Working with lookup tables diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index 59c6998..eb714e5 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -1729,26 +1729,26 @@ class Table(Queryable): columns[0], fn, drop=drop, show_progress=show_progress ) - todo_count = self.count * len(columns) - if output is not None: + assert len(columns) == 1, "output= can only be used with a single column" if output not in self.columns_dict: self.add_column(output, output_type or "text") + todo_count = self.count * len(columns) with progressbar(length=todo_count, silent=not show_progress) as bar: - def transform_value(v): + def convert_value(v): bar.update(1) if not v: return v return fn(v) - self.db.register_function(transform_value) + self.db.register_function(convert_value) sql = "update [{table}] set {sets};".format( table=self.name, sets=", ".join( [ - "[{output_column}] = transform_value([{column}])".format( + "[{output_column}] = convert_value([{column}])".format( output_column=output or column, column=column ) for column in columns diff --git a/tests/test_convert.py b/tests/test_convert.py index 15e8026..34e98f1 100644 --- a/tests/test_convert.py +++ b/tests/test_convert.py @@ -38,6 +38,13 @@ def test_convert_output(fresh_db, drop, expected): assert list(table.rows) == [expected] +def test_convert_output_multiple_column_error(fresh_db): + table = fresh_db["table"] + with pytest.raises(AssertionError) as excinfo: + table.convert(["title", "other"], lambda v: v, output="out") + assert "output= can only be used with a single column" in str(excinfo.value) + + @pytest.mark.parametrize( "type,expected", ( From 652040f21313f4eda4cb63485b9b8e2668314db8 Mon Sep 17 00:00:00 2001 From: Simon Willison Date: Sun, 1 Aug 2021 22:05:03 -0700 Subject: [PATCH 18/18] table.count_where() method, closes #305 --- docs/python-api.rst | 12 +++++++++++- sqlite_utils/db.py | 32 +++++++++++++++++++++++++------- tests/test_enable_counts.py | 2 +- tests/test_introspect.py | 9 ++++++++- 4 files changed, 45 insertions(+), 10 deletions(-) diff --git a/docs/python-api.rst b/docs/python-api.rst index 7f490ab..aba43ad 100644 --- a/docs/python-api.rst +++ b/docs/python-api.rst @@ -286,6 +286,16 @@ This method also accepts ``offset=`` and ``limit=`` arguments, for specifying an ... print(row) {'id': 1, 'age': 4, 'name': 'Cleo'} +.. _python_api_rows_count_where: + +Counting rows +------------- + +To count the number of rows that would be returned by a where filter, use ``.count_where(where, where_args)``: + + >>> db["dogs"].count_where("age > ?", [1]): + 2 + .. _python_api_pks_and_rows_where: Listing rows with their primary keys @@ -1602,7 +1612,7 @@ The ``.count`` property shows the current number of rows (``select count(*) from >>> db["Street_Tree_List"].count 189144 -This property will take advantage of :ref:`python_api_cached_table_counts` if the ``use_counts_table`` property is set on the database. You can avoid that optimization entirely by calling ``table.execute_count()`` instead of accessing the property. +This property will take advantage of :ref:`python_api_cached_table_counts` if the ``use_counts_table`` property is set on the database. You can avoid that optimization entirely by calling ``table.count_where()`` instead of accessing the property. .. _python_api_introspection_columns: diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index eb714e5..b5da058 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -682,14 +682,23 @@ class Queryable: self.db = db self.name = name + def count_where( + self, + where=None, + where_args=None, + ): + sql = "select count(*) from [{}]".format(self.name) + if where is not None: + sql += " where " + where + return self.db.execute(sql, where_args or []).fetchone()[0] + def execute_count(self): - return self.db.execute( - "select count(*) from [{}]".format(self.name) - ).fetchone()[0] + # Backwards compatibility, see https://github.com/simonw/sqlite-utils/issues/305#issuecomment-890713185 + return self.count_where() @property def count(self): - return self.execute_count() + return self.count_where() @property def rows(self): @@ -820,7 +829,7 @@ class Table(Queryable): counts = self.db.cached_counts([self.name]) if counts: return next(iter(counts.values())) - return self.execute_count() + return self.count_where() def exists(self): return self.name in self.db.table_names() @@ -1719,6 +1728,8 @@ class Table(Queryable): output_type=None, drop=False, multi=False, + where=None, + where_args=None, show_progress=False, ): if isinstance(columns, str): @@ -1726,7 +1737,12 @@ class Table(Queryable): if multi: return self._convert_multi( - columns[0], fn, drop=drop, show_progress=show_progress + columns[0], + fn, + drop=drop, + where=where, + where_args=where_args, + show_progress=show_progress, ) if output is not None: @@ -1761,7 +1777,9 @@ class Table(Queryable): self.transform(drop=columns) return self - def _convert_multi(self, column, fn, drop, show_progress): + def _convert_multi( + self, column, fn, drop, show_progress, where=None, where_args=None + ): # First we execute the function pk_to_values = {} new_column_types = {} diff --git a/tests/test_enable_counts.py b/tests/test_enable_counts.py index 7a52108..d724e80 100644 --- a/tests/test_enable_counts.py +++ b/tests/test_enable_counts.py @@ -132,7 +132,7 @@ def test_uses_counts_after_enable_counts(counts_db_path): assert db["foo"].count == 1 assert logged == [ ("select name from sqlite_master where type = 'view'", None), - ("select count(*) from [foo]", None), + ("select count(*) from [foo]", []), ] logged.clear() assert not db.use_counts_table diff --git a/tests/test_introspect.py b/tests/test_introspect.py index cc33c46..dce8afc 100644 --- a/tests/test_introspect.py +++ b/tests/test_introspect.py @@ -52,7 +52,14 @@ def test_views(fresh_db): def test_count(existing_db): - assert 3 == existing_db["foo"].count + assert existing_db["foo"].count == 3 + assert existing_db["foo"].count_where() == 3 + assert existing_db["foo"].execute_count() == 3 + + +def test_count_where(existing_db): + assert existing_db["foo"].count_where("text != ?", ["two"]) == 2 + assert existing_db["foo"].count_where("text != :t", {"t": "two"}) == 2 def test_columns(existing_db):