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]