Fix all remaining resource warnings, refs #693

https://gistpreview.github.io/?0bb8e869b82f6ff0db647de755182502

Author: simonw

sqlite_utils/cli.py

@@ -905,7 +905,7 @@ def inner(fn):
                     required=True,
                 ),
                 click.argument("table"),
-                click.argument("file", type=click.File("rb"), required=True),
+                click.argument("file", type=click.File("rb", lazy=True), required=True),
                 click.option(
                     "--pk",
                     help="Columns to use as the primary key, e.g. id",
@@ -2000,6 +2000,7 @@ def memory(
     for i, path in enumerate(paths):
         # Path may have a :format suffix
         fp = None
+        should_close_fp = False
         if ":" in path and path.rsplit(":", 1)[-1].upper() in Format.__members__:
             path, suffix = path.rsplit(":", 1)
             format = Format[suffix.upper()]
@@ -2017,29 +2018,32 @@ def memory(
                 file_table = stem
             stem_counts[stem] = stem_counts.get(stem, 1) + 1
             fp = file_path.open("rb")
-        rows, format_used = rows_from_file(fp, format=format, encoding=encoding)
-        tracker = None
-        if format_used in (Format.CSV, Format.TSV) and not no_detect_types:
-            tracker = TypeTracker()
-            rows = tracker.wrap(rows)
-        if flatten:
-            rows = (_flatten(row) for row in rows)
-
-        db[file_table].insert_all(rows, alter=True)
-        if tracker is not None:
-            db[file_table].transform(types=tracker.types)
-        # Add convenient t / t1 / t2 views
-        view_names = ["t{}".format(i + 1)]
-        if i == 0:
-            view_names.append("t")
-        for view_name in view_names:
-            if not db[view_name].exists():
-                db.create_view(
-                    view_name, "select * from {}".format(quote_identifier(file_table))
-                )
-
-        if fp:
-            fp.close()
+            should_close_fp = True
+        try:
+            rows, format_used = rows_from_file(fp, format=format, encoding=encoding)
+            tracker = None
+            if format_used in (Format.CSV, Format.TSV) and not no_detect_types:
+                tracker = TypeTracker()
+                rows = tracker.wrap(rows)
+            if flatten:
+                rows = (_flatten(row) for row in rows)
+
+            db[file_table].insert_all(rows, alter=True)
+            if tracker is not None:
+                db[file_table].transform(types=tracker.types)
+            # Add convenient t / t1 / t2 views
+            view_names = ["t{}".format(i + 1)]
+            if i == 0:
+                view_names.append("t")
+            for view_name in view_names:
+                if not db[view_name].exists():
+                    db.create_view(
+                        view_name,
+                        "select * from {}".format(quote_identifier(file_table)),
+                    )
+        finally:
+            if should_close_fp and fp:
+                fp.close()
 
     if analyze:
         _analyze(db, tables=None, columns=None, save=False)

sqlite_utils/utils.py

@@ -9,7 +9,29 @@
 import os
 import sys
 from . import recipes
-from typing import Dict, cast, BinaryIO, Iterable, Optional, Tuple, Type
+from typing import Dict, cast, BinaryIO, Iterable, Iterator, Optional, Tuple, Type
+
+
+class _CloseableIterator(Iterator[dict]):
+    """Iterator wrapper that closes a file when iteration is complete."""
+
+    def __init__(self, iterator: Iterator[dict], closeable: io.IOBase):
+        self._iterator = iterator
+        self._closeable = closeable
+
+    def __iter__(self) -> "_CloseableIterator":
+        return self
+
+    def __next__(self) -> dict:
+        try:
+            return next(self._iterator)
+        except StopIteration:
+            self._closeable.close()
+            raise
+
+    def close(self) -> None:
+        self._closeable.close()
+
 
 import click
 
@@ -299,7 +321,8 @@ class Format(enum.Enum):
             reader = csv.DictReader(decoded_fp, dialect=dialect)
         else:
             reader = csv.DictReader(decoded_fp)
-        return _extra_key_strategy(reader, ignore_extras, extras_key), Format.CSV
+        rows = _extra_key_strategy(reader, ignore_extras, extras_key)
+        return _CloseableIterator(iter(rows), decoded_fp), Format.CSV
     elif format == Format.TSV:
         rows = rows_from_file(
             fp, format=Format.CSV, dialect=csv.excel_tab, encoding=encoding

tests/test_cli.py

@@ -19,9 +19,11 @@ def _supports_pragma_function_list():
     db = Database(memory=True)
     try:
         db.execute("select * from pragma_function_list()")
+        return True
     except Exception:
         return False
-    return True
+    finally:
+        db.close()
 
 
 def _has_compiled_ext():

tests/test_introspect.py

@@ -2,6 +2,14 @@
 import pytest
 
 
+def _check_supports_strict():
+    """Check if SQLite supports strict tables without leaking the database."""
+    db = Database(memory=True)
+    result = db.supports_strict
+    db.close()
+    return result
+
+
 def test_table_names(existing_db):
     assert ["foo"] == existing_db.table_names()
 
@@ -282,7 +290,7 @@ def test_use_rowid(fresh_db):
 
 
 @pytest.mark.skipif(
-    not Database(memory=True).supports_strict,
+    not _check_supports_strict(),
     reason="Needs SQLite version that supports strict",
 )
 @pytest.mark.parametrize(

tests/test_plugins.py

@@ -9,9 +9,11 @@ def _supports_pragma_function_list():
     db = Database(memory=True)
     try:
         db.execute("select * from pragma_function_list()")
+        return True
     except Exception:
         return False
-    return True
+    finally:
+        db.close()
 
 
 def test_register_commands():

tests/test_recreate.py

@@ -14,8 +14,11 @@ def test_recreate_ignored_for_in_memory():
 
 def test_recreate_not_allowed_for_connection():
     conn = sqlite3.connect(":memory:")
-    with pytest.raises(AssertionError):
-        Database(conn, recreate=True)
+    try:
+        with pytest.raises(AssertionError):
+            Database(conn, recreate=True)
+    finally:
+        conn.close()
 
 
 @pytest.mark.parametrize(

tests/test_register_function.py

@@ -42,39 +42,46 @@ def to_lower(s):
 
 
 def test_register_function_deterministic_tries_again_if_exception_raised(fresh_db):
+    # Save the original connection so we can close it later
+    original_conn = fresh_db.conn
     fresh_db.conn = MagicMock()
     fresh_db.conn.create_function = MagicMock()
 
-    @fresh_db.register_function(deterministic=True)
-    def to_lower_2(s):
-        return s.lower()
-
-    fresh_db.conn.create_function.assert_called_with(
-        "to_lower_2", 1, to_lower_2, deterministic=True
-    )
-
-    first = True
-
-    def side_effect(*args, **kwargs):
-        # Raise exception only first time this is called
-        nonlocal first
-        if first:
-            first = False
-            raise sqlite3.NotSupportedError()
-
-    # But if sqlite3.NotSupportedError is raised, it tries again
-    fresh_db.conn.create_function.reset_mock()
-    fresh_db.conn.create_function.side_effect = side_effect
-
-    @fresh_db.register_function(deterministic=True)
-    def to_lower_3(s):
-        return s.lower()
-
-    # Should have been called once with deterministic=True and once without
-    assert fresh_db.conn.create_function.call_args_list == [
-        call("to_lower_3", 1, to_lower_3, deterministic=True),
-        call("to_lower_3", 1, to_lower_3),
-    ]
+    try:
+
+        @fresh_db.register_function(deterministic=True)
+        def to_lower_2(s):
+            return s.lower()
+
+        fresh_db.conn.create_function.assert_called_with(
+            "to_lower_2", 1, to_lower_2, deterministic=True
+        )
+
+        first = True
+
+        def side_effect(*args, **kwargs):
+            # Raise exception only first time this is called
+            nonlocal first
+            if first:
+                first = False
+                raise sqlite3.NotSupportedError()
+
+        # But if sqlite3.NotSupportedError is raised, it tries again
+        fresh_db.conn.create_function.reset_mock()
+        fresh_db.conn.create_function.side_effect = side_effect
+
+        @fresh_db.register_function(deterministic=True)
+        def to_lower_3(s):
+            return s.lower()
+
+        # Should have been called once with deterministic=True and once without
+        assert fresh_db.conn.create_function.call_args_list == [
+            call("to_lower_3", 1, to_lower_3, deterministic=True),
+            call("to_lower_3", 1, to_lower_3),
+        ]
+    finally:
+        # Close the original connection that was replaced with the mock
+        original_conn.close()
 
 
 def test_register_function_replace(fresh_db):