From 2f38660d48b1dc55b2ed5b30e86cc7c5b04494be Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Tue, 27 Oct 2020 11:07:06 -0700
Subject: [PATCH] Progress bar for insert works now

---
 sqlite_utils/utils.py | 21 ++++++++-------------
 tests/test_cli.py     |  6 ++++--
 2 files changed, 12 insertions(+), 15 deletions(-)

diff --git a/sqlite_utils/utils.py b/sqlite_utils/utils.py
index aab72ca..a158b2b 100644
--- a/sqlite_utils/utils.py
+++ b/sqlite_utils/utils.py
@@ -92,20 +92,15 @@ def find_spatialite():
     return None
 
 
-class UpdateReader(io.TextIOWrapper):
-    def __init__(self, raw, update):
-        super().__init__(raw)
+class UpdateWrapper:
+    def __init__(self, wrapped, update):
+        self._wrapped = wrapped
         self._update = update
 
-    def read(self, size=-1):
-        bytes = super().read(size)
-        self._update(len(bytes))
-        return bytes
-
-    def readline(self, size=-1):
-        bytes = super().readline(size)
-        self._update(len(bytes))
-        return bytes
+    def __iter__(self):
+        for line in self._wrapped:
+            self._update(len(line))
+            yield line
 
 
 @contextlib.contextmanager
@@ -115,4 +110,4 @@ def file_progress(file, silent=False, **kwargs):
     else:
         file_length = os.path.getsize(file.raw.name)
         with click.progressbar(length=file_length, **kwargs) as bar:
-            yield UpdateReader(file, update=bar.update)
+            yield UpdateWrapper(file, bar.update)
diff --git a/tests/test_cli.py b/tests/test_cli.py
index 5c3056e..7e1365d 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -1640,7 +1640,9 @@ def test_insert_encoding(tmpdir):
     open(csv_path, "wb").write(latin1_csv)
     # First attempt should error:
     bad_result = CliRunner().invoke(
-        cli.cli, ["insert", db_path, "places", csv_path, "--csv"], catch_exceptions=False
+        cli.cli,
+        ["insert", db_path, "places", csv_path, "--csv"],
+        catch_exceptions=False,
     )
     assert bad_result.exit_code == 1
     assert (
@@ -1651,7 +1653,7 @@ def test_insert_encoding(tmpdir):
     good_result = CliRunner().invoke(
         cli.cli,
         ["insert", db_path, "places", csv_path, "--encoding", "latin-1", "--csv"],
-        catch_exceptions=False
+        catch_exceptions=False,
     )
     assert good_result.exit_code == 0
     db = Database(db_path)