From 3e04646fd4a1bb64f6cd5a1cd3058e603b5cb1ae Mon Sep 17 00:00:00 2001
From: collerek <collerek@gmail.com>
Date: Tue, 11 Aug 2020 18:32:48 +0200
Subject: [PATCH] refactors in fk

---
 .coverage                 | Bin 53248 -> 53248 bytes
 orm/fields/foreign_key.py |  45 +++++++++++++++++++-------------------
 2 files changed, 23 insertions(+), 22 deletions(-)

diff --git a/.coverage b/.coverage
index 105fceb53d64948024e1f0c6aae03edf3748d6b8..12fa709d6622e39bed47bdc45522a017c60c3638 100644
GIT binary patch
delta 93
zcmV-j0HXhZpaX!Q1F$tO1u{80H8`_6FU(LRTL2IF59$xz57Q6G53>)S4~-8~4+9SE
z4&@Hl4z3QA4ss4#vk?$64zpa2L;(~M2Lu5LCI{NZpX`5k_uD=0+>RBqvyV+c%eo=u

delta 92
zcmV-i0HgnapaX!Q1F$tO1u!@|Ffg+^FU(LRTmTRG59$xz57Q6G53>)T4~`F14+IYG
y4(1Nn4zCWC4s#A%vk?$74zpX1L;({K2Lu5LBnR5XpX`5k_qJQdj!v_)k4-?Mtsx8m

diff --git a/orm/fields/foreign_key.py b/orm/fields/foreign_key.py
index b5659b4..827ab19 100644
--- a/orm/fields/foreign_key.py
+++ b/orm/fields/foreign_key.py
@@ -25,12 +25,12 @@ def create_dummy_instance(fk: Type["Model"], pk: int = None) -> "Model":
 
 class ForeignKey(BaseField):
     def __init__(
-        self,
-        to: Type["Model"],
-        name: str = None,
-        related_name: str = None,
-        nullable: bool = True,
-        virtual: bool = False,
+            self,
+            to: Type["Model"],
+            name: str = None,
+            related_name: str = None,
+            nullable: bool = True,
+            virtual: bool = False,
     ) -> None:
         super().__init__(nullable=nullable, name=name)
         self.virtual = virtual
@@ -50,36 +50,42 @@ class ForeignKey(BaseField):
         return to_column.get_column_type()
 
     def extract_model_from_sequence(
-        self, value: Any, child: "Model"
-    ) -> Tuple[Union["Model", List["Model"]], bool]:
+            self, value: Any, child: "Model"
+    ) -> Union["Model", List["Model"]]:
         if isinstance(value, list) and not isinstance(value, self.to):
             model = [self.expand_relationship(val, child) for val in value]
-            return model, True
+            return model
 
         if isinstance(value, self.to):
             model = value
         else:
             model = self.to(**value)
-        return model, False
+        self.register_relation(model, child)
+        return model
 
-    def construct_model_from_pk(self, value: Any) -> "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__} "
                 f"is of type {self.to.pk_type()} "
                 f"while {type(value)} passed as a parameter."
             )
-        return create_dummy_instance(fk=self.to, pk=value)
+        model = create_dummy_instance(fk=self.to, pk=value)
+        self.register_relation(model, child)
+        return model
+
+    def register_relation(self, model, child):
+        model._orm_relationship_manager.add_relation(
+            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
 
-        is_sequence = False
-
         if isinstance(value, orm.models.Model) and not isinstance(value, self.to):
             raise RelationshipInstanceError(
                 f"Relationship error - expecting: {self.to.__name__}, "
@@ -87,13 +93,8 @@ class ForeignKey(BaseField):
             )
 
         if isinstance(value, (dict, list, self.to)):
-            model, is_sequence = self.extract_model_from_sequence(value, child)
+            model = self.extract_model_from_sequence(value, child)
         else:
-            model = self.construct_model_from_pk(value)
-
-        if not is_sequence:
-            model._orm_relationship_manager.add_relation(
-                model, child, virtual=self.virtual
-            )
+            model = self.construct_model_from_pk(value, child)
 
         return model