Address review feedback for Annotated relationships

- Detect annotated relationships without a default value by iterating
  over the union of class_dict and original_annotations.
- Use elif and unwrap inner Mapped[T] when handling Annotated[Mapped[T], ...]
  to avoid double-wrapping in Mapped.
- Split the relationship/pydantic partition into three single-purpose
  loops with comments for clarity.
- Add tests covering Annotated relationships with default value, without
  default value, and with Annotated[Mapped[T], ...].

Author: bolu61

sqlmodel/main.py

@@ -561,20 +561,29 @@ def __new__(
         original_annotations = get_annotations(class_dict)
         pydantic_annotations = {}
         relationship_annotations = {}
-        for k, v in class_dict.items():
-            a = original_annotations.get(k, None)
-            r = get_annotated_relationshipinfo(a)
+
+        # find relationship info in both annotations and class dict
+        for k in {**original_annotations, **class_dict}:
+            v = class_dict.get(k)
+            if isinstance(v, RelationshipInfo):
+                relationships[k] = v
+                continue
+            r = get_annotated_relationshipinfo(original_annotations.get(k))
             if r is not None:
                 relationships[k] = r
-            elif isinstance(v, RelationshipInfo):
-                relationships[k] = v
-            else:
+
+        # populate dict passed to pydantic
+        for k, v in class_dict.items():
+            if k not in relationships:
                 dict_for_pydantic[k] = v
-        for k, v in original_annotations.items():
+
+        # split out pydantic annotations
+        for k, a in original_annotations.items():
             if k in relationships:
-                relationship_annotations[k] = v
+                relationship_annotations[k] = a
             else:
-                pydantic_annotations[k] = v
+                pydantic_annotations[k] = a
+
         dict_used = {
             **dict_for_pydantic,
             "__weakref__": None,
@@ -659,8 +668,10 @@ def __init__(
                 origin: Any = get_origin(raw_ann)
                 if origin is Mapped:
                     ann = raw_ann.__args__[0]
-                if origin is Annotated:
+                elif origin is Annotated:
                     ann = get_args(raw_ann)[0]
+                    if get_origin(ann) is Mapped:
+                        ann = ann.__args__[0]
                     cls.__annotations__[rel_name] = Mapped[ann]  # type: ignore[valid-type]
                 else:
                     ann = raw_ann

tests/test_annotated_relationship.py

@@ -0,0 +1,85 @@
+from typing import Annotated
+
+from sqlalchemy.orm import Mapped
+from sqlmodel import Field, Relationship, Session, SQLModel, create_engine, select
+
+
+def test_annotated_relationship_with_default() -> None:
+    class Team(SQLModel, table=True):
+        id: Annotated[int | None, Field(primary_key=True)] = None
+        name: Annotated[str, Field(index=True)]
+
+        heroes: Annotated[list["Hero"], Relationship(back_populates="team")] = []  # noqa: RUF012
+
+    class Hero(SQLModel, table=True):
+        id: Annotated[int | None, Field(primary_key=True)] = None
+        name: Annotated[str, Field(index=True)]
+        team_id: Annotated[int | None, Field(foreign_key="team.id")] = None
+        team: Annotated[Team | None, Relationship(back_populates="heroes")] = None
+
+    engine = create_engine("sqlite://")
+    SQLModel.metadata.create_all(engine)
+    with Session(engine) as session:
+        team = Team(name="Preventers")
+        hero = Hero(name="Deadpond", team=team)
+        session.add(hero)
+        session.commit()
+        session.refresh(hero)
+        assert hero.team is not None
+        assert hero.team.name == "Preventers"
+        team_db = session.exec(select(Team)).one()
+        assert [h.name for h in team_db.heroes] == ["Deadpond"]
+
+
+def test_annotated_relationship_without_default() -> None:
+    class Team(SQLModel, table=True):
+        id: Annotated[int | None, Field(primary_key=True)] = None
+        name: Annotated[str, Field(index=True)]
+
+        heroes: Annotated[list["Hero"], Relationship(back_populates="team")]
+
+    class Hero(SQLModel, table=True):
+        id: Annotated[int | None, Field(primary_key=True)] = None
+        name: Annotated[str, Field(index=True)]
+        team_id: Annotated[int | None, Field(foreign_key="team.id")] = None
+        team: Annotated[Team | None, Relationship(back_populates="heroes")]
+
+    engine = create_engine("sqlite://")
+    SQLModel.metadata.create_all(engine)
+    with Session(engine) as session:
+        team = Team(name="Z-Force")
+        hero = Hero(name="Spider-Boy", team=team)
+        session.add(hero)
+        session.commit()
+        session.refresh(hero)
+        assert hero.team is not None
+        assert hero.team.name == "Z-Force"
+
+
+def test_annotated_mapped_relationship() -> None:
+    class Team(SQLModel, table=True):
+        id: Annotated[int | None, Field(primary_key=True)] = None
+        name: Annotated[str, Field(index=True)]
+
+        heroes: Annotated[
+            Mapped[list["Hero"]], Relationship(back_populates="team")
+        ] = []  # noqa: RUF012
+
+    class Hero(SQLModel, table=True):
+        id: Annotated[int | None, Field(primary_key=True)] = None
+        name: Annotated[str, Field(index=True)]
+        team_id: Annotated[int | None, Field(foreign_key="team.id")] = None
+        team: Annotated[Mapped[Team | None], Relationship(back_populates="heroes")] = (
+            None
+        )
+
+    engine = create_engine("sqlite://")
+    SQLModel.metadata.create_all(engine)
+    with Session(engine) as session:
+        team = Team(name="Avengers")
+        hero = Hero(name="Iron Man", team=team)
+        session.add(hero)
+        session.commit()
+        session.refresh(hero)
+        assert hero.team is not None
+        assert hero.team.name == "Avengers"