From 975b6c6191c82d1e62100032d335bd03832775f9 Mon Sep 17 00:00:00 2001 From: Simon Willison Date: Sun, 23 Nov 2025 09:29:43 -0800 Subject: [PATCH] Support tuples as well as lists Refs #672 --- sqlite_utils/db.py | 24 +++++++---- tests/test_list_mode.py | 90 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 106 insertions(+), 8 deletions(-) diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index a8ed704..e8cfe9e 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -3033,14 +3033,18 @@ class Table(Queryable): # Pad short records with None, truncate long ones record_len = len(record) if record_len < num_columns: - record_values = [jsonify_if_needed(v) for v in record] + [None] * (num_columns - record_len) + record_values = [jsonify_if_needed(v) for v in record] + [None] * ( + num_columns - record_len + ) else: record_values = [jsonify_if_needed(v) for v in record[:num_columns]] # Only process extracts if there are any if has_extracts: for i, key in enumerate(all_columns): if key in extracts: - record_values[i] = self.db[extracts[key]].lookup({"value": record_values[i]}) + record_values[i] = self.db[extracts[key]].lookup( + {"value": record_values[i]} + ) values.append(record_values) else: # Dict mode: original logic @@ -3393,12 +3397,14 @@ class Table(Queryable): return self # It was an empty list # Check if this is list mode or dict mode - if isinstance(first_record, list): - # List mode: first record should be column names + if isinstance(first_record, (list, tuple)): + # List/tuple 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 + raise ValueError( + "When using list-based iteration, the first yielded value must be a list of column name strings" + ) + column_names = list(first_record) all_columns = column_names num_columns = len(column_names) # Get the actual first data record @@ -3406,8 +3412,10 @@ class Table(Queryable): 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") + if not isinstance(first_record, (list, tuple)): + 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) diff --git a/tests/test_list_mode.py b/tests/test_list_mode.py index 4f841ff..01f0a49 100644 --- a/tests/test_list_mode.py +++ b/tests/test_list_mode.py @@ -1,6 +1,7 @@ """ Tests for list-based iteration in insert_all and upsert_all """ + import pytest from sqlite_utils import Database @@ -173,3 +174,92 @@ def test_backwards_compatibility_dict_mode(): rows = list(db["people"].rows) assert len(rows) == 2 assert rows[0] == {"id": 1, "name": "Alice", "age": 30} + + +def test_insert_all_tuple_mode_basic(): + """Test basic insert_all with tuple-based iteration""" + db = Database(memory=True) + + def data_generator(): + # First yield column names as tuple + yield ("id", "name", "age") + # Then yield data rows as tuples + 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_mixed_list_tuple(): + """Test insert_all with mixed lists and tuples for data rows""" + db = Database(memory=True) + + def data_generator(): + # Column names as list + yield ["id", "name", "age"] + # Mix of list and tuple data rows + yield [1, "Alice", 30] + yield (2, "Bob", 25) + yield [3, "Charlie", 35] + yield (4, "Diana", 40) + + db["people"].insert_all(data_generator()) + + rows = list(db["people"].rows) + assert len(rows) == 4 + 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} + assert rows[3] == {"id": 4, "name": "Diana", "age": 40} + + +def test_upsert_all_tuple_mode(): + """Test upsert_all with tuple-based iteration""" + db = Database(memory=True) + + # Initial insert with tuples + 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 tuples + 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_tuple_mode_shorter_rows(): + """Test that tuple 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}