refactor expanding of relationship into constructors

orm/fields/foreign_key.py

@@ -1,4 +1,4 @@
-from typing import Any, List, Optional, TYPE_CHECKING, Tuple, Type, Union
+from typing import Any, List, Optional, TYPE_CHECKING, Type, Union
 
 import sqlalchemy
 from pydantic import BaseModel
@@ -49,21 +49,21 @@ class ForeignKey(BaseField):
         to_column = self.to.__model_fields__[self.to.__pkname__]
         return to_column.get_column_type()
 
-    def extract_model_from_sequence(
+    def _extract_model_from_sequence(
-            self, value: Any, child: "Model"
+        self, value: List, child: "Model"
     ) -> Union["Model", List["Model"]]:
-        if isinstance(value, list) and not isinstance(value, self.to):
+        return [self.expand_relationship(val, child) for val in value]
-            model = [self.expand_relationship(val, child) for val in value]
-            return model
 
-        if isinstance(value, self.to):
+    def _register_existing_model(self, value: "Model", child: "Model") -> "Model":
-            model = value
+        self.register_relation(value, child)
-        else:
+        return value
+
+    def _construct_model_from_dict(self, value: dict, child: "Model") -> "Model":
         model = self.to(**value)
         self.register_relation(model, child)
         return model
 
-    def construct_model_from_pk(self, value: Any, child: "Model") -> "Model":
+    def _construct_model_from_pk(self, value: Any, child: "Model") -> "Model":
         if not isinstance(value, self.to.pk_type()):
             raise RelationshipInstanceError(
                 f"Relationship error - ForeignKey {self.to.__name__} "
@@ -74,10 +74,8 @@ class ForeignKey(BaseField):
         self.register_relation(model, child)
         return model
 
-    def register_relation(self, model, child):
+    def register_relation(self, model: "Model", child: "Model") -> None:
-        model._orm_relationship_manager.add_relation(
+        model._orm_relationship_manager.add_relation(model, child, virtual=self.virtual)
-            model, child, virtual=self.virtual
-        )
 
     def expand_relationship(
         self, value: Any, child: "Model"
@@ -86,15 +84,13 @@ class ForeignKey(BaseField):
         if value is None:
             return None
 
-        if isinstance(value, orm.models.Model) and not isinstance(value, self.to):
+        constructors = {
-            raise RelationshipInstanceError(
+            f"{self.to.__name__}": self._register_existing_model,
-                f"Relationship error - expecting: {self.to.__name__}, "
+            "dict": self._construct_model_from_dict,
-                f"but {value.__class__.__name__} encountered."
+            "list": self._extract_model_from_sequence,
-            )
+        }
-
-        if isinstance(value, (dict, list, self.to)):
-            model = self.extract_model_from_sequence(value, child)
-        else:
-            model = self.construct_model_from_pk(value, child)
 
+        model = constructors.get(
+            value.__class__.__name__, self._construct_model_from_pk
+        )(value, child)
         return model