diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index af5bb97..0a16380 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -1,5 +1,6 @@ from .utils import sqlite3, OperationalError, suggest_column_types, column_affinity from collections import namedtuple, OrderedDict +from collections.abc import Mapping import contextlib import datetime import decimal @@ -1772,15 +1773,15 @@ class Table(Queryable): return pk def m2m( - self, other_table, record_or_list=None, pk=DEFAULT, lookup=None, m2m_table=None + self, other_table, record_or_iterable=None, pk=DEFAULT, lookup=None, m2m_table=None ): if isinstance(other_table, str): other_table = self.db.table(other_table, pk=pk) our_id = self.last_pk if lookup is not None: - assert record_or_list is None, "Provide lookup= or record, not both" + assert record_or_iterable is None, "Provide lookup= or record, not both" else: - assert record_or_list is not None, "Provide lookup= or record, not both" + assert record_or_iterable is not None, "Provide lookup= or record, not both" tables = list(sorted([self.name, other_table.name])) columns = ["{}_id".format(t) for t in tables] if m2m_table is not None: @@ -1801,10 +1802,11 @@ class Table(Queryable): m2m_table_name = m2m_table or "{}_{}".format(*tables) m2m_table = self.db.table(m2m_table_name, pk=columns, foreign_keys=columns) if lookup is None: + # if records is only one record, put the record in a list records = ( - [record_or_list] - if not isinstance(record_or_list, (list, tuple)) - else record_or_list + [record_or_iterable] + if isinstance(record_or_iterable, Mapping) + else record_or_iterable ) # Ensure each record exists in other table for record in records: diff --git a/tests/test_m2m.py b/tests/test_m2m.py index 52b3959..33cd0a8 100644 --- a/tests/test_m2m.py +++ b/tests/test_m2m.py @@ -31,16 +31,52 @@ def test_insert_m2m_list(fresh_db): humans.rows ) assert [ - ForeignKey( - table="dogs_humans", column="dogs_id", other_table="dogs", other_column="id" - ), - ForeignKey( - table="dogs_humans", - column="humans_id", - other_table="humans", - other_column="id", - ), - ] == dogs_humans.foreign_keys + ForeignKey( + table="dogs_humans", column="dogs_id", other_table="dogs", other_column="id" + ), + ForeignKey( + table="dogs_humans", + column="humans_id", + other_table="humans", + other_column="id", + ), + ] == dogs_humans.foreign_keys + + +def test_insert_m2m_iterable(fresh_db): + iterable_records = ({"id": 1, "name": "Phineas"}, {"id": 2, "name": "Ferb"}) + + def iterable(): + for record in iterable_records: + yield record + + platypuses = fresh_db["platypuses"] + platypuses.insert({"id": 1, "name": "Perry"}, pk="id").m2m( + "humans", + iterable(), + pk="id", + ) + + assert {"platypuses", "humans", "humans_platypuses"} == set(fresh_db.table_names()) + humans = fresh_db["humans"] + humans_platypuses = fresh_db["humans_platypuses"] + assert [{"humans_id": 1, "platypuses_id": 1}, {"humans_id": 2, "platypuses_id": 1}] == list( + humans_platypuses.rows + ) + assert [{"id": 1, "name": "Phineas"}, {"id": 2, "name": "Ferb"}] == list( + humans.rows + ) + assert [ + ForeignKey( + table="humans_platypuses", column="platypuses_id", other_table="platypuses", other_column="id" + ), + ForeignKey( + table="humans_platypuses", + column="humans_id", + other_table="humans", + other_column="id", + ), + ] == humans_platypuses.foreign_keys def test_m2m_with_table_objects(fresh_db): @@ -64,16 +100,16 @@ def test_m2m_lookup(fresh_db): assert people_tags.exists() assert tags.exists() assert [ - ForeignKey( - table="people_tags", - column="people_id", - other_table="people", - other_column="id", - ), - ForeignKey( - table="people_tags", column="tags_id", other_table="tags", other_column="id" - ), - ] == people_tags.foreign_keys + ForeignKey( + table="people_tags", + column="people_id", + other_table="people", + other_column="id", + ), + ForeignKey( + table="people_tags", column="tags_id", other_table="tags", other_column="id" + ), + ] == people_tags.foreign_keys assert [{"people_id": 1, "tags_id": 1}] == list(people_tags.rows) assert [{"id": 1, "name": "Wahyu"}] == list(people.rows) assert [{"id": 1, "tag": "Coworker"}] == list(tags.rows)