Initial Claude implementation of list mode

From https://github.com/simonw/research/tree/main/sqlite-utils-iterator-support

I made very minor changes to the test (removing a sys.path insert)

Refs #672

Refs https://github.com/simonw/research/pull/31

sqlite_utils/db.py

@@ -3010,6 +3010,7 @@ class Table(Queryable):
         num_records_processed,
         replace,
         ignore,
+        list_mode=False,
     ):
         """
         Given a list ``chunk`` of records that should be written to *this* table,
@@ -3024,24 +3025,40 @@ class Table(Queryable):
         # Build a row-list ready for executemany-style flattening
         values = []
 
-        for record in chunk:
-            record_values = []
-            for key in all_columns:
-                value = jsonify_if_needed(
-                    record.get(
-                        key,
-                        (
-                            None
-                            if key != hash_id
-                            else hash_record(record, hash_id_columns)
-                        ),
+        if list_mode:
+            # In list mode, records are already lists of values
+            for record in chunk:
+                record_values = []
+                for i, key in enumerate(all_columns):
+                    if i < len(record):
+                        value = jsonify_if_needed(record[i])
+                    else:
+                        value = None
+                    if key in extracts:
+                        extract_table = extracts[key]
+                        value = self.db[extract_table].lookup({"value": value})
+                    record_values.append(value)
+                values.append(record_values)
+        else:
+            # Dict mode: original logic
+            for record in chunk:
+                record_values = []
+                for key in all_columns:
+                    value = jsonify_if_needed(
+                        record.get(
+                            key,
+                            (
+                                None
+                                if key != hash_id
+                                else hash_record(record, hash_id_columns)
+                            ),
+                        )
                     )
-                )
-                if key in extracts:
-                    extract_table = extracts[key]
-                    value = self.db[extract_table].lookup({"value": value})
-                record_values.append(value)
-            values.append(record_values)
+                    if key in extracts:
+                        extract_table = extracts[key]
+                        value = self.db[extract_table].lookup({"value": value})
+                    record_values.append(value)
+                values.append(record_values)
 
         columns_sql = ", ".join(f"[{c}]" for c in all_columns)
         placeholder_expr = ", ".join(conversions.get(c, "?") for c in all_columns)
@@ -3157,6 +3174,7 @@ class Table(Queryable):
         num_records_processed,
         replace,
         ignore,
+        list_mode=False,
     ) -> Optional[sqlite3.Cursor]:
         queries_and_params = self.build_insert_queries_and_params(
             extracts,
@@ -3171,6 +3189,7 @@ class Table(Queryable):
             num_records_processed,
             replace,
             ignore,
+            list_mode,
         )
         result = None
         with self.db.conn:
@@ -3200,6 +3219,7 @@ class Table(Queryable):
                             num_records_processed,
                             replace,
                             ignore,
+                            list_mode,
                         )
 
                         result = self.insert_chunk(
@@ -3216,6 +3236,7 @@ class Table(Queryable):
                             num_records_processed,
                             replace,
                             ignore,
+                            list_mode,
                         )
 
                     else:
@@ -3353,17 +3374,47 @@ class Table(Queryable):
         all_columns = []
         first = True
         num_records_processed = 0
-        # Fix up any records with square braces in the column names
-        records = fix_square_braces(records)
-        # We can only handle a max of 999 variables in a SQL insert, so
-        # we need to adjust the batch_size down if we have too many cols
-        records = iter(records)
-        # Peek at first record to count its columns:
+
+        # Detect if we're using list-based iteration or dict-based iteration
+        list_mode = False
+        column_names = None
+
+        # Fix up any records with square braces in the column names (only for dict mode)
+        # We'll handle this differently for list mode
+        records_iter = iter(records)
+
+        # Peek at first record to determine mode:
         try:
-            first_record = next(records)
+            first_record = next(records_iter)
         except StopIteration:
             return self  # It was an empty list
-        num_columns = len(first_record.keys())
+
+        # Check if this is list mode or dict mode
+        if isinstance(first_record, list):
+            # List mode: first record should be column names
+            list_mode = True
+            if not all(isinstance(col, str) for col in first_record):
+                raise ValueError("When using list-based iteration, the first yielded value must be a list of column name strings")
+            column_names = first_record
+            all_columns = column_names
+            num_columns = len(column_names)
+            # Get the actual first data record
+            try:
+                first_record = next(records_iter)
+            except StopIteration:
+                return self  # Only headers, no data
+            if not isinstance(first_record, list):
+                raise ValueError("After column names list, all subsequent records must also be lists")
+        else:
+            # Dict mode: traditional behavior
+            records_iter = itertools.chain([first_record], records_iter)
+            records_iter = fix_square_braces(records_iter)
+            try:
+                first_record = next(records_iter)
+            except StopIteration:
+                return self
+            num_columns = len(first_record.keys())
+
         assert (
             num_columns <= SQLITE_MAX_VARS
         ), "Rows can have a maximum of {} columns".format(SQLITE_MAX_VARS)
@@ -3373,13 +3424,18 @@ class Table(Queryable):
         if truncate and self.exists():
             self.db.execute("DELETE FROM [{}];".format(self.name))
         result = None
-        for chunk in chunks(itertools.chain([first_record], records), batch_size):
+        for chunk in chunks(itertools.chain([first_record], records_iter), batch_size):
             chunk = list(chunk)
             num_records_processed += len(chunk)
             if first:
                 if not self.exists():
                     # Use the first batch to derive the table names
-                    column_types = suggest_column_types(chunk)
+                    if list_mode:
+                        # Convert list records to dicts for type detection
+                        chunk_as_dicts = [dict(zip(column_names, row)) for row in chunk]
+                        column_types = suggest_column_types(chunk_as_dicts)
+                    else:
+                        column_types = suggest_column_types(chunk)
                     if extracts:
                         for col in extracts:
                             if col in column_types:
@@ -3399,17 +3455,24 @@ class Table(Queryable):
                         extracts=extracts,
                         strict=strict,
                     )
-                all_columns_set = set()
-                for record in chunk:
-                    all_columns_set.update(record.keys())
-                all_columns = list(sorted(all_columns_set))
-                if hash_id:
-                    all_columns.insert(0, hash_id)
+                if list_mode:
+                    # In list mode, columns are already known
+                    all_columns = list(column_names)
+                    if hash_id:
+                        all_columns.insert(0, hash_id)
+                else:
+                    all_columns_set = set()
+                    for record in chunk:
+                        all_columns_set.update(record.keys())
+                    all_columns = list(sorted(all_columns_set))
+                    if hash_id:
+                        all_columns.insert(0, hash_id)
             else:
-                for record in chunk:
-                    all_columns += [
-                        column for column in record if column not in all_columns
-                    ]
+                if not list_mode:
+                    for record in chunk:
+                        all_columns += [
+                            column for column in record if column not in all_columns
+                        ]
 
             first = False
 
@@ -3427,6 +3490,7 @@ class Table(Queryable):
                 num_records_processed,
                 replace,
                 ignore,
+                list_mode,
             )
 
         # If we only handled a single row populate self.last_pk

tests/test_list_mode.py

@@ -0,0 +1,175 @@
+"""
+Tests for list-based iteration in insert_all and upsert_all
+"""
+import pytest
+from sqlite_utils import Database
+
+
+def test_insert_all_list_mode_basic():
+    """Test basic insert_all with list-based iteration"""
+    db = Database(memory=True)
+
+    def data_generator():
+        # First yield column names
+        yield ["id", "name", "age"]
+        # Then yield data rows
+        yield [1, "Alice", 30]
+        yield [2, "Bob", 25]
+        yield [3, "Charlie", 35]
+
+    db["people"].insert_all(data_generator())
+
+    rows = list(db["people"].rows)
+    assert len(rows) == 3
+    assert rows[0] == {"id": 1, "name": "Alice", "age": 30}
+    assert rows[1] == {"id": 2, "name": "Bob", "age": 25}
+    assert rows[2] == {"id": 3, "name": "Charlie", "age": 35}
+
+
+def test_insert_all_list_mode_with_pk():
+    """Test insert_all with list mode and primary key"""
+    db = Database(memory=True)
+
+    def data_generator():
+        yield ["id", "name", "score"]
+        yield [1, "Alice", 95]
+        yield [2, "Bob", 87]
+
+    db["scores"].insert_all(data_generator(), pk="id")
+
+    assert db["scores"].pks == ["id"]
+    rows = list(db["scores"].rows)
+    assert len(rows) == 2
+
+
+def test_upsert_all_list_mode():
+    """Test upsert_all with list-based iteration"""
+    db = Database(memory=True)
+
+    # Initial insert
+    def initial_data():
+        yield ["id", "name", "value"]
+        yield [1, "Alice", 100]
+        yield [2, "Bob", 200]
+
+    db["data"].insert_all(initial_data(), pk="id")
+
+    # Upsert with some updates and new records
+    def upsert_data():
+        yield ["id", "name", "value"]
+        yield [1, "Alice", 150]  # Update existing
+        yield [3, "Charlie", 300]  # Insert new
+
+    db["data"].upsert_all(upsert_data(), pk="id")
+
+    rows = list(db["data"].rows_where(order_by="id"))
+    assert len(rows) == 3
+    assert rows[0] == {"id": 1, "name": "Alice", "value": 150}
+    assert rows[1] == {"id": 2, "name": "Bob", "value": 200}
+    assert rows[2] == {"id": 3, "name": "Charlie", "value": 300}
+
+
+def test_list_mode_with_various_types():
+    """Test list mode with different data types"""
+    db = Database(memory=True)
+
+    def data_generator():
+        yield ["id", "name", "score", "active"]
+        yield [1, "Alice", 95.5, True]
+        yield [2, "Bob", 87.3, False]
+        yield [3, "Charlie", None, True]
+
+    db["mixed"].insert_all(data_generator())
+
+    rows = list(db["mixed"].rows)
+    assert len(rows) == 3
+    assert rows[0]["score"] == 95.5
+    assert rows[1]["active"] == 0  # SQLite stores boolean as int
+    assert rows[2]["score"] is None
+
+
+def test_list_mode_error_non_string_columns():
+    """Test that non-string column names raise an error"""
+    db = Database(memory=True)
+
+    def bad_data():
+        yield [1, 2, 3]  # Non-string column names
+        yield ["a", "b", "c"]
+
+    with pytest.raises(ValueError, match="must be a list of column name strings"):
+        db["bad"].insert_all(bad_data())
+
+
+def test_list_mode_error_mixed_types():
+    """Test that mixing list and dict raises an error"""
+    db = Database(memory=True)
+
+    def bad_data():
+        yield ["id", "name"]
+        yield {"id": 1, "name": "Alice"}  # Should be a list, not dict
+
+    with pytest.raises(ValueError, match="must also be lists"):
+        db["bad"].insert_all(bad_data())
+
+
+def test_list_mode_empty_after_headers():
+    """Test that only headers without data works gracefully"""
+    db = Database(memory=True)
+
+    def data_generator():
+        yield ["id", "name", "age"]
+        # No data rows
+
+    result = db["people"].insert_all(data_generator())
+    assert result is not None
+    assert not db["people"].exists()
+
+
+def test_list_mode_batch_processing():
+    """Test list mode with large dataset requiring batching"""
+    db = Database(memory=True)
+
+    def large_data():
+        yield ["id", "value"]
+        for i in range(1000):
+            yield [i, f"value_{i}"]
+
+    db["large"].insert_all(large_data(), batch_size=100)
+
+    count = db.execute("SELECT COUNT(*) as c FROM large").fetchone()[0]
+    assert count == 1000
+
+
+def test_list_mode_shorter_rows():
+    """Test that rows shorter than column list get NULL values"""
+    db = Database(memory=True)
+
+    def data_generator():
+        yield ["id", "name", "age", "city"]
+        yield [1, "Alice", 30, "NYC"]
+        yield [2, "Bob"]  # Missing age and city
+        yield [3, "Charlie", 35]  # Missing city
+
+    db["people"].insert_all(data_generator())
+
+    rows = list(db["people"].rows_where(order_by="id"))
+    assert rows[0] == {"id": 1, "name": "Alice", "age": 30, "city": "NYC"}
+    assert rows[1] == {"id": 2, "name": "Bob", "age": None, "city": None}
+    assert rows[2] == {"id": 3, "name": "Charlie", "age": 35, "city": None}
+
+
+def test_backwards_compatibility_dict_mode():
+    """Ensure dict mode still works (backward compatibility)"""
+    db = Database(memory=True)
+
+    # Traditional dict-based insert
+    data = [
+        {"id": 1, "name": "Alice", "age": 30},
+        {"id": 2, "name": "Bob", "age": 25},
+    ]
+
+    db["people"].insert_all(data)
+
+    rows = list(db["people"].rows)
+    assert len(rows) == 2
+    assert rows[0] == {"id": 1, "name": "Alice", "age": 30}