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
@@ -25,12 +25,12 @@ def create_dummy_instance(fk: Type["Model"], pk: int = None) -> "Model":
 
 class ForeignKey(BaseField):
     def __init__(
-            self,
+        self,
-            to: Type["Model"],
+        to: Type["Model"],
-            name: str = None,
+        name: str = None,
-            related_name: str = None,
+        related_name: str = None,
-            nullable: bool = True,
+        nullable: bool = True,
-            virtual: bool = False,
+        virtual: bool = False,
     ) -> None:
         super().__init__(nullable=nullable, name=name)
         self.virtual = virtual
@@ -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
-            model = self.to(**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,27 +74,23 @@ 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"
+        self, value: Any, child: "Model"
     ) -> Optional[Union["Model", List["Model"]]]:
 
         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