From 570bee7edd17424e37ce7d7e6a1b35d4be063f17 Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Sun, 1 Aug 2021 21:08:40 -0700
Subject: [PATCH] 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: