table.convert(...) method, refs #251

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,

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]