From b47d5209eb4a21787bdfd30e374de19730eb0f3b Mon Sep 17 00:00:00 2001
From: collerek <collerek@gmail.com>
Date: Tue, 15 Sep 2020 13:53:54 +0200
Subject: [PATCH] some refactors in relations module

---
 .coverage                           | Bin 53248 -> 53248 bytes
 ormar/models/modelproxy.py          |  17 ++++++++++++++++-
 ormar/relations/relation.py         |   2 +-
 ormar/relations/relation_manager.py |  16 +---------------
 4 files changed, 18 insertions(+), 17 deletions(-)

diff --git a/.coverage b/.coverage
index acd8ef5ebb1d924c99fdf1ab0bae137f94dc01d4..428308a7ce8aaf4d69734fa00147102962e93359 100644
GIT binary patch
delta 78
zcmV-U0I~mopaX!Q1F$eJ1T`=^F|#x;`%nR2voL@c4lWiI1OW*y751gc`)_~$_P@RF
ku@C+BdoL6e1OW*u6!we%y`R1JzyI&u{dv28vvrP(LALxO=Kufz

delta 78
zcmV-U0I~mopaX!Q1F$eJ1Tr-`GP5)<`%nR0voL@c4lWcG1OW*u6?QRs|LyPJ{<rr%
k_S^5h78C>l2`&`&i~qf!z4yQW@7-U5KX3OBvvrP(L9AaR0{{R3

diff --git a/ormar/models/modelproxy.py b/ormar/models/modelproxy.py
index 81e1ea6..2cd86a2 100644
--- a/ormar/models/modelproxy.py
+++ b/ormar/models/modelproxy.py
@@ -1,13 +1,17 @@
 import inspect
-from typing import List, Optional, Set, TYPE_CHECKING
+from typing import List, Optional, Set, TYPE_CHECKING, Type, TypeVar
 
 import ormar
+from ormar.exceptions import RelationshipInstanceError
+from ormar.fields import BaseField
 from ormar.fields.foreign_key import ForeignKeyField
 from ormar.models.metaclass import ModelMeta
 
 if TYPE_CHECKING:  # pragma no cover
     from ormar import Model
 
+Field = TypeVar("Field", bound=BaseField)
+
 
 class ModelTableProxy:
     if TYPE_CHECKING:  # pragma no cover
@@ -89,6 +93,17 @@ class ModelTableProxy:
                 if field.to == related.__class__ or field.to.Meta == related.Meta:
                     return name
 
+    @staticmethod
+    def resolve_relation_field(item: "Model", related: "Model") -> Type[Field]:
+        name = ModelTableProxy.resolve_relation_name(item, related)
+        to_field = item.Meta.model_fields.get(name)
+        if not to_field:  # pragma no cover
+            raise RelationshipInstanceError(
+                f"Model {item.__class__} does not have "
+                f"reference to model {related.__class__}"
+            )
+        return to_field
+
     @classmethod
     def merge_instances_list(cls, result_rows: List["Model"]) -> List["Model"]:
         merged_rows = []
diff --git a/ormar/relations/relation.py b/ormar/relations/relation.py
index 01e0ec4..e0747fb 100644
--- a/ormar/relations/relation.py
+++ b/ormar/relations/relation.py
@@ -39,7 +39,7 @@ class Relation:
     def _find_existing(self, child: "Model") -> Optional[int]:
         for ind, relation_child in enumerate(self.related_models[:]):
             try:
-                if relation_child.__same__(child):
+                if relation_child == child:
                     return ind
             except ReferenceError:  # pragma no cover
                 self.related_models.pop(ind)
diff --git a/ormar/relations/relation_manager.py b/ormar/relations/relation_manager.py
index 1407182..9574374 100644
--- a/ormar/relations/relation_manager.py
+++ b/ormar/relations/relation_manager.py
@@ -2,7 +2,6 @@ from typing import List, Optional, TYPE_CHECKING, Tuple, Type, Union
 from weakref import proxy
 
 import ormar
-from ormar.exceptions import RelationshipInstanceError
 from ormar.fields.foreign_key import ForeignKeyField
 from ormar.fields.many_to_many import ManyToManyField
 from ormar.relations import Relation
@@ -85,20 +84,7 @@ class RelationsManager:
 
     @staticmethod
     def add(parent: "Model", child: "Model", child_name: str, virtual: bool) -> None:
-        to_field = next(
-            (
-                field
-                for field in child._orm._related_fields
-                if field.to == parent.__class__ or field.to.Meta == parent.Meta
-            ),
-            None,
-        )
-
-        if not to_field:  # pragma no cover
-            raise RelationshipInstanceError(
-                f"Model {child.__class__} does not have "
-                f"reference to model {parent.__class__}"
-            )
+        to_field = child.resolve_relation_field(child, parent)
 
         (
             parent,