Incorporate review feedback

Author: Natarajan Krishnaswami

sqlmodel/_compat.py

@@ -178,16 +178,16 @@ def get_relationship_to(
         # If a list, then also get the real field
         elif origin is list:
             use_annotation = get_args(annotation)[0]
-        # If a dict or Mapping, then use the value (second) type argument
-        elif origin is dict or origin is Mapping:
+        # If a dict, then use the value (second) type argument
+        elif origin is dict:
             args = get_args(annotation)
-            if len(args) >= 2:
-                use_annotation = args[1]
-            else:
+            if len(args) != 2:
                 raise ValueError(
-                    f"Dict/Mapping relationship field '{name}' must have both "
-                    "key and value type arguments (e.g., dict[str, Model])"
+                    f"Dict/Mapping relationship field '{name}' has {len(args)} "
+                    "type arguments.  Exactly two required (e.g., dict[str, "
+                    "Model])"
                 )
+            use_annotation = args[1]
 
         return get_relationship_to(
             name=name, rel_info=rel_info, annotation=use_annotation

tests/test_attribute_keyed_dict.py

@@ -7,7 +7,7 @@
 from sqlalchemy.orm.collections import attribute_keyed_dict
 from sqlmodel import Field, Relationship, Session, SQLModel, create_engine
 
-from tests.conftest import needs_pydanticv2
+from tests.conftest import needs_py310, needs_pydanticv2
 
 
 def test_attribute_keyed_dict_works(clear_sqlmodel):
@@ -53,8 +53,8 @@ class Parent(SQLModel, table=True):
 # typing.Dict throws if it receives the wrong number of type arguments, but dict
 # (3.10+) does not; and Pydantic v1 fails to process models with dicts with no
 # type arguments.
-@pytest.mark.skipif(sys.version_info < (3, 10), reason="dict is not subscriptable")
 @needs_pydanticv2
+@needs_py310
 def test_dict_relationship_throws_on_missing_annotation_arg(clear_sqlmodel):
     class Color(str, Enum):
         Orange = "Orange"
@@ -68,11 +68,10 @@ class Child(SQLModel, table=True):
         color: Color
         value: int
 
-    error_msg_re = re.escape(
-        "Dict/Mapping relationship field 'children_by_color' must have both key and value type arguments (e.g., dict[str, Model])"
-    )
+    error_msg_fmt =  "Dict/Mapping relationship field 'children_by_color' has {count} type arguments.  Exactly two required (e.g., dict[str, Model])"
+
     # No type args
-    with pytest.raises(ValueError, match=error_msg_re):
+    with pytest.raises(ValueError, match=re.escape(error_msg_fmt.format(count=0))):
 
         class Parent(SQLModel, table=True):
             __tablename__ = "parents"
@@ -85,7 +84,7 @@ class Parent(SQLModel, table=True):
             )
 
     # One type arg
-    with pytest.raises(ValueError, match=error_msg_re):
+    with pytest.raises(ValueError, match=re.escape(error_msg_fmt.format(count=1))):
 
         class Parent(SQLModel, table=True):
             __tablename__ = "parents"
@@ -96,3 +95,16 @@ class Parent(SQLModel, table=True):
                     "collection_class": attribute_keyed_dict("color")
                 }
             )
+
+    # Three type args
+    with pytest.raises(ValueError, match=re.escape(error_msg_fmt.format(count=3))):
+
+        class Parent(SQLModel, table=True):
+            __tablename__ = "parents"
+
+            id: Optional[int] = Field(primary_key=True, default=None)
+            children_by_color: dict[Color, Child, str] = Relationship(
+                sa_relationship_kwargs={
+                    "collection_class": attribute_keyed_dict("color")
+                }
+            )