Better type signatures for compound foreign key columns

New ForeignKeyColumns type alias - a column name or a list/tuple of
column names - used by add_foreign_key(), add_foreign_keys() (via the
ForeignKeyTuple alias, replacing its Any slots) and the
ForeignKeyIndicator union, which it also simplifies. Tuples now
type-check anywhere lists are accepted.

ForeignKey.__post_init__ normalizes tuple columns/other_columns to
lists so direct construction with tuples compares equal to
introspected foreign keys.

Refs https://github.com/simonw/sqlite-utils/pull/770/changes#r3525703477

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Simon Willison 2026-07-05 15:03:23 -07:00
commit c16edb2dc4
2 changed files with 36 additions and 13 deletions

View file

@ -197,10 +197,15 @@ class ForeignKey:
on_update: str = "NO ACTION" on_update: str = "NO ACTION"
def __post_init__(self): def __post_init__(self):
# Populate columns/other_columns for single-column foreign keys # Populate columns/other_columns for single-column foreign keys,
if not self.columns: # normalizing any tuples to lists
if self.columns:
self.columns = list(self.columns)
else:
self.columns = [self.column] if self.column is not None else [] self.columns = [self.column] if self.column is not None else []
if not self.other_columns: if self.other_columns:
self.other_columns = list(self.other_columns)
else:
self.other_columns = ( self.other_columns = (
[self.other_column] if self.other_column is not None else [] [self.other_column] if self.other_column is not None else []
) )
@ -228,16 +233,18 @@ class TransformError(Exception):
pass pass
# A single column name, or a list/tuple of columns for a compound foreign key
ForeignKeyColumns = Union[str, List[str], Tuple[str, ...]]
# (table, column(s), other_table, other_column(s))
ForeignKeyTuple = Tuple[str, ForeignKeyColumns, str, ForeignKeyColumns]
ForeignKeyIndicator = Union[ ForeignKeyIndicator = Union[
str, str,
ForeignKey, ForeignKey,
Tuple[str, str], Tuple[ForeignKeyColumns, str],
Tuple[str, str, str], Tuple[ForeignKeyColumns, str, ForeignKeyColumns],
Tuple[str, str, str, str], ForeignKeyTuple,
# Compound foreign keys use lists of columns:
Tuple[List[str], str],
Tuple[List[str], str, List[str]],
Tuple[str, List[str], str, List[str]],
] ]
ForeignKeysType = Union[Iterable[ForeignKeyIndicator], List[ForeignKeyIndicator]] ForeignKeysType = Union[Iterable[ForeignKeyIndicator], List[ForeignKeyIndicator]]
@ -1604,7 +1611,7 @@ class Database:
return candidates return candidates
def add_foreign_keys( def add_foreign_keys(
self, foreign_keys: Iterable[Union[ForeignKey, Tuple[str, Any, str, Any]]] self, foreign_keys: Iterable[Union[ForeignKey, ForeignKeyTuple]]
) -> None: ) -> None:
""" """
See :ref:`python_api_add_foreign_keys`. See :ref:`python_api_add_foreign_keys`.
@ -2911,9 +2918,9 @@ class Table(Queryable):
def add_foreign_key( def add_foreign_key(
self, self,
column: Union[str, List[str]], column: ForeignKeyColumns,
other_table: Optional[str] = None, other_table: Optional[str] = None,
other_column: Optional[Union[str, List[str]]] = None, other_column: Optional[ForeignKeyColumns] = None,
ignore: bool = False, ignore: bool = False,
): ):
""" """

View file

@ -445,3 +445,19 @@ def test_implicit_compound_primary_key_reference_is_resolved():
fk = db["courses"].foreign_keys[0] fk = db["courses"].foreign_keys[0]
assert fk.is_compound is True assert fk.is_compound is True
assert fk.other_columns == ["campus_name", "dept_code"] assert fk.other_columns == ["campus_name", "dept_code"]
def test_foreign_key_normalizes_tuple_columns_to_lists():
# Compound columns passed as tuples are normalized to lists, so they
# compare equal to introspected ForeignKeys
fk = ForeignKey(
table="courses",
column=None,
other_table="departments",
other_column=None,
columns=("campus_name", "dept_code"),
other_columns=("campus_name", "dept_code"),
is_compound=True,
)
assert fk.columns == ["campus_name", "dept_code"]
assert fk.other_columns == ["campus_name", "dept_code"]