From 69c7da5ec9698dabeb23379cc08d012b0cd8e6d2 Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Mon, 2 Aug 2021 11:33:56 -0700
Subject: [PATCH] Implemented .convert(..., where=, where_args=), refs #304

---
 docs/python-api.rst   |  8 ++++++++
 sqlite_utils/db.py    | 11 +++++++----
 tests/test_convert.py | 40 ++++++++++++++++++++++++++++++++++++++++
 3 files changed, 55 insertions(+), 4 deletions(-)

diff --git a/docs/python-api.rst b/docs/python-api.rst
index aba43ad..c372a62 100644
--- a/docs/python-api.rst
+++ b/docs/python-api.rst
@@ -785,6 +785,14 @@ You can create multiple new columns from a single input column by passing ``mult
         "title", lambda v: {"upper": v.upper(), "lower": v.lower()}, multi=True
     )
 
+The ``.convert()`` method accepts optional ``where=`` and ``where_args=`` parameters which can be used to apply the conversion to a subset of rows specified by a where clause. Here's how to apply the conversion only to rows with an ``id`` that is higher than 20:
+
+.. code-block:: python
+
+    table.convert("title", lambda v: v.upper(), where="id > :id", where_args={"id": 20})
+
+These behave the same as the corresponding parameters to the :ref:`.rows_where() <python_api_rows>` method, so you can use ``?`` placeholders and a list of values instead of ``:named`` placeholders with a dictionary.
+
 .. _python_api_lookup_tables:
 
 Working with lookup tables
diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py
index b5da058..a21acc4 100644
--- a/sqlite_utils/db.py
+++ b/sqlite_utils/db.py
@@ -1750,7 +1750,7 @@ class Table(Queryable):
             if output not in self.columns_dict:
                 self.add_column(output, output_type or "text")
 
-        todo_count = self.count * len(columns)
+        todo_count = self.count_where(where, where_args) * len(columns)
         with progressbar(length=todo_count, silent=not show_progress) as bar:
 
             def convert_value(v):
@@ -1760,7 +1760,7 @@ class Table(Queryable):
                 return fn(v)
 
             self.db.register_function(convert_value)
-            sql = "update [{table}] set {sets};".format(
+            sql = "update [{table}] set {sets}{where};".format(
                 table=self.name,
                 sets=", ".join(
                     [
@@ -1770,9 +1770,10 @@ class Table(Queryable):
                         for column in columns
                     ]
                 ),
+                where=" where {}".format(where) if where is not None else "",
             )
             with self.db.conn:
-                self.db.execute(sql)
+                self.db.execute(sql, where_args or [])
                 if drop:
                     self.transform(drop=columns)
         return self
@@ -1793,7 +1794,9 @@ class Table(Queryable):
             for row in self.rows_where(
                 select=", ".join(
                     "[{}]".format(column_name) for column_name in (pks + [column])
-                )
+                ),
+                where=where,
+                where_args=where_args,
             ):
                 row_pk = tuple(row[pk] for pk in pks)
                 if len(row_pk) == 1:
diff --git a/tests/test_convert.py b/tests/test_convert.py
index 34e98f1..796a08f 100644
--- a/tests/test_convert.py
+++ b/tests/test_convert.py
@@ -24,6 +24,24 @@ def test_convert(fresh_db, columns, fn, expected):
     assert list(table.rows) == [expected]
 
 
+@pytest.mark.parametrize(
+    "where,where_args", (("id > 1", None), ("id > :id", {"id": 1}), ("id > ?", [1]))
+)
+def test_convert_where(fresh_db, where, where_args):
+    table = fresh_db["table"]
+    table.insert_all(
+        [
+            {"id": 1, "title": "One"},
+            {"id": 2, "title": "Two"},
+        ],
+        pk="id",
+    )
+    table.convert(
+        "title", lambda value: value.upper(), where=where, where_args=where_args
+    )
+    assert list(table.rows) == [{"id": 1, "title": "One"}, {"id": 2, "title": "TWO"}]
+
+
 @pytest.mark.parametrize(
     "drop,expected",
     (
@@ -70,6 +88,28 @@ def test_convert_multi(fresh_db):
     ]
 
 
+def test_convert_multi_where(fresh_db):
+    table = fresh_db["table"]
+    table.insert_all(
+        [
+            {"id": 1, "title": "One"},
+            {"id": 2, "title": "Two"},
+        ],
+        pk="id",
+    )
+    table.convert(
+        "title",
+        lambda v: {"upper": v.upper(), "lower": v.lower()},
+        multi=True,
+        where="id > ?",
+        where_args=[1],
+    )
+    assert list(table.rows) == [
+        {"id": 1, "lower": None, "title": "One", "upper": None},
+        {"id": 2, "lower": "two", "title": "Two", "upper": "TWO"},
+    ]
+
+
 def test_convert_multi_exception(fresh_db):
     table = fresh_db["table"]
     table.insert({"title": "Mixed Case"})