From e8f21c56242978d32d50df2841e35898a3172c06 Mon Sep 17 00:00:00 2001
From: Mark Pullin <mark.pullin@exponential.eco>
Date: Thu, 4 May 2023 00:19:24 +0100
Subject: [PATCH] refactor

---
 ormar/models/helpers/validation.py    |  4 +---
 ormar/models/mixins/merge_mixin.py    |  4 +---
 ormar/models/mixins/relation_mixin.py | 11 ++++++++++-
 ormar/models/model.py                 |  4 +---
 ormar/models/newbasemodel.py          | 14 +-------------
 5 files changed, 14 insertions(+), 23 deletions(-)

diff --git a/ormar/models/helpers/validation.py b/ormar/models/helpers/validation.py
index d25837dc1..eae1afd87 100644
--- a/ormar/models/helpers/validation.py
+++ b/ormar/models/helpers/validation.py
@@ -108,9 +108,7 @@ def generate_model_example(model: Type["Model"], relation_map: Dict = None) -> D
     """
     example: Dict[str, Any] = dict()
     relation_map = (
-        relation_map
-        if relation_map is not None
-        else translate_list_to_dict(model._iterate_related_models())
+        relation_map if relation_map is not None else model._related_models_dict()
     )
     for name, field in model.Meta.model_fields.items():
         populates_sample_fields_values(
diff --git a/ormar/models/mixins/merge_mixin.py b/ormar/models/mixins/merge_mixin.py
index 7ef20778c..d39451855 100644
--- a/ormar/models/mixins/merge_mixin.py
+++ b/ormar/models/mixins/merge_mixin.py
@@ -87,9 +87,7 @@ def merge_two_instances(
         :rtype: Model
         """
         relation_map = (
-            relation_map
-            if relation_map is not None
-            else translate_list_to_dict(one._iterate_related_models())
+            relation_map if relation_map is not None else one._related_models_dict()
         )
         for field_name in relation_map:
             current_field = getattr(one, field_name)
diff --git a/ormar/models/mixins/relation_mixin.py b/ormar/models/mixins/relation_mixin.py
index 3f9521c5c..a491bcdae 100644
--- a/ormar/models/mixins/relation_mixin.py
+++ b/ormar/models/mixins/relation_mixin.py
@@ -1,7 +1,8 @@
-from typing import Callable, Dict, List, Optional, Set, TYPE_CHECKING, cast
+from typing import Any, Callable, Dict, List, Optional, Set, TYPE_CHECKING, cast
 
 from ormar import BaseField, ForeignKeyField
 from ormar.models.traversible import NodeList
+from ormar.queryset.utils import translate_list_to_dict
 
 
 class RelationMixin:
@@ -14,6 +15,7 @@ class RelationMixin:
 
         Meta: ModelMeta
         __relation_map__: Optional[List[str]]
+        __relation_map_dict__: Optional[Dict[str, Any]]
         _related_names: Optional[Set]
         _through_names: Optional[Set]
         _related_fields: Optional[List]
@@ -112,6 +114,13 @@ def _extract_db_related_names(cls) -> Set:
         }
         return related_names
 
+    @classmethod
+    def _related_models_dict(cls) -> dict[str, Any]:
+        if not cls.__relation_map_dict__:
+            cls.__relation_map_dict__ = translate_list_to_dict(cls._iterate_related_models())
+
+        return cls.__relation_map_dict__
+
     @classmethod
     def _iterate_related_models(  # noqa: CCR001
         cls,
diff --git a/ormar/models/model.py b/ormar/models/model.py
index 7ff8e6369..f74f53d17 100644
--- a/ormar/models/model.py
+++ b/ormar/models/model.py
@@ -153,9 +153,7 @@ async def save_related(  # noqa: CCR001, CFQ002
         :rtype: int
         """
         relation_map = (
-            relation_map
-            if relation_map is not None
-            else translate_list_to_dict(self._iterate_related_models())
+            relation_map if relation_map is not None else self._related_models_dict()
         )
         if exclude and isinstance(exclude, Set):
             exclude = translate_list_to_dict(exclude)
diff --git a/ormar/models/newbasemodel.py b/ormar/models/newbasemodel.py
index 4790c9be6..8357d36b6 100644
--- a/ormar/models/newbasemodel.py
+++ b/ormar/models/newbasemodel.py
@@ -735,16 +735,6 @@ def _extract_nested_models(  # noqa: CCR001, CFQ002
             except ReferenceError:
                 dict_instance[field] = None
         return dict_instance
-    
-
-    @classmethod
-    def _related_models_dict(cls) -> dict:
-        if not cls.__relation_map_dict__:
-            cls.__relation_map_dict__ = translate_list_to_dict(
-                cls._iterate_related_models()
-                )
-        return cls.__relation_map_dict__
-        
 
     def dict(  # type: ignore # noqa A003
         self,
@@ -822,9 +812,7 @@ def dict(  # type: ignore # noqa A003
             exclude = translate_list_to_dict(exclude)
 
         relation_map = (
-            relation_map
-            if relation_map is not None
-            else self._related_models_dict()
+            relation_map if relation_map is not None else self._related_models_dict()
         )
         pk_only = getattr(self, "__pk_only__", False)
         if relation_map and not pk_only: