refactor and cleanup - drop of resolving relation names as not fully proper, extract mixins from modelproxy to be more maintainable, add some docstrings

ormar/models/mixins/__init__.py

@@ -0,0 +1,5 @@
+from ormar.models.mixins.alias_mixin import AliasMixin
+from ormar.models.mixins.merge_mixin import MergeModelMixin
+from ormar.models.mixins.prefetch_mixin import PrefetchQueryMixin
+
+__all__ = ["MergeModelMixin", "AliasMixin", "PrefetchQueryMixin"]

ormar/models/mixins/alias_mixin.py

@@ -0,0 +1,83 @@
+from typing import Dict, List, Optional, Set, TYPE_CHECKING, Type, Union
+
+
+class AliasMixin:
+    if TYPE_CHECKING:  # pragma: no cover
+        from ormar import Model, ModelMeta
+
+        Meta: ModelMeta
+
+    @classmethod
+    def get_column_alias(cls, field_name: str) -> str:
+        field = cls.Meta.model_fields.get(field_name)
+        return field.get_alias() if field is not None else field_name
+
+    @classmethod
+    def get_column_name_from_alias(cls, alias: str) -> str:
+        for field_name, field in cls.Meta.model_fields.items():
+            if field.get_alias() == alias:
+                return field_name
+        return alias  # if not found it's not an alias but actual name
+
+    @classmethod
+    def translate_columns_to_aliases(cls, new_kwargs: Dict) -> Dict:
+        for field_name, field in cls.Meta.model_fields.items():
+            if field_name in new_kwargs:
+                new_kwargs[field.get_alias()] = new_kwargs.pop(field_name)
+        return new_kwargs
+
+    @classmethod
+    def translate_aliases_to_columns(cls, new_kwargs: Dict) -> Dict:
+        for field_name, field in cls.Meta.model_fields.items():
+            if field.alias and field.alias in new_kwargs:
+                new_kwargs[field_name] = new_kwargs.pop(field.alias)
+        return new_kwargs
+
+    @staticmethod
+    def _populate_pk_column(
+        model: Type["Model"], columns: List[str], use_alias: bool = False,
+    ) -> List[str]:
+        pk_alias = (
+            model.get_column_alias(model.Meta.pkname)
+            if use_alias
+            else model.Meta.pkname
+        )
+        if pk_alias not in columns:
+            columns.append(pk_alias)
+        return columns
+
+    @classmethod
+    def own_table_columns(
+        cls,
+        model: Type["Model"],
+        fields: Optional[Union[Set, Dict]],
+        exclude_fields: Optional[Union[Set, Dict]],
+        use_alias: bool = False,
+    ) -> List[str]:
+        columns = [
+            model.get_column_name_from_alias(col.name) if not use_alias else col.name
+            for col in model.Meta.table.columns
+        ]
+        field_names = [
+            model.get_column_name_from_alias(col.name)
+            for col in model.Meta.table.columns
+        ]
+        if fields:
+            columns = [
+                col
+                for col, name in zip(columns, field_names)
+                if model.is_included(fields, name)
+            ]
+        if exclude_fields:
+            columns = [
+                col
+                for col, name in zip(columns, field_names)
+                if not model.is_excluded(exclude_fields, name)
+            ]
+
+        # always has to return pk column for ormar to work
+        columns = cls._populate_pk_column(
+            model=model, columns=columns, use_alias=use_alias
+        )
+
+        return columns

ormar/models/mixins/merge_mixin.py

@@ -0,0 +1,46 @@
+from collections import OrderedDict
+from typing import List, Sequence, TYPE_CHECKING
+
+import ormar
+
+if TYPE_CHECKING:  # pragma no cover
+    from ormar import Model
+
+
+class MergeModelMixin:
+    @classmethod
+    def merge_instances_list(cls, result_rows: Sequence["Model"]) -> Sequence["Model"]:
+        merged_rows: List["Model"] = []
+        grouped_instances: OrderedDict = OrderedDict()
+
+        for model in result_rows:
+            grouped_instances.setdefault(model.pk, []).append(model)
+
+        for group in grouped_instances.values():
+            model = group.pop(0)
+            if group:
+                for next_model in group:
+                    model = cls.merge_two_instances(next_model, model)
+            merged_rows.append(model)
+
+        return merged_rows
+
+    @classmethod
+    def merge_two_instances(cls, one: "Model", other: "Model") -> "Model":
+        for field in one.Meta.model_fields.keys():
+            current_field = getattr(one, field)
+            if isinstance(current_field, list) and not isinstance(
+                current_field, ormar.Model
+            ):
+                setattr(other, field, current_field + getattr(other, field))
+            elif (
+                isinstance(current_field, ormar.Model)
+                and current_field.pk == getattr(other, field).pk
+            ):
+                setattr(
+                    other,
+                    field,
+                    cls.merge_two_instances(current_field, getattr(other, field)),
+                )
+        other.set_save_status(True)
+        return other

ormar/models/mixins/prefetch_mixin.py

@@ -0,0 +1,64 @@
+from typing import Callable, Dict, List, TYPE_CHECKING, Tuple, Type
+
+import ormar
+from ormar.fields import BaseField
+
+
+class PrefetchQueryMixin:
+    if TYPE_CHECKING:  # pragma no cover
+        from ormar import Model
+
+        get_name: Callable  # defined in NewBaseModel
+        extract_related_names: Callable  # defined in ModelTableProxy
+
+    @staticmethod
+    def get_clause_target_and_filter_column_name(
+        parent_model: Type["Model"],
+        target_model: Type["Model"],
+        reverse: bool,
+        related: str,
+    ) -> Tuple[Type["Model"], str]:
+        if reverse:
+            field_name = (
+                parent_model.Meta.model_fields[related].related_name
+                or parent_model.get_name() + "s"
+            )
+            field = target_model.Meta.model_fields[field_name]
+            if issubclass(field, ormar.fields.ManyToManyField):
+                field_name = field.default_target_field_name()
+                sub_field = field.through.Meta.model_fields[field_name]
+                return field.through, sub_field.get_alias()
+            return target_model, field.get_alias()
+        target_field = target_model.get_column_alias(target_model.Meta.pkname)
+        return target_model, target_field
+
+    @staticmethod
+    def get_column_name_for_id_extraction(
+        parent_model: Type["Model"], reverse: bool, related: str, use_raw: bool,
+    ) -> str:
+        if reverse:
+            column_name = parent_model.Meta.pkname
+            return (
+                parent_model.get_column_alias(column_name) if use_raw else column_name
+            )
+        column = parent_model.Meta.model_fields[related]
+        return column.get_alias() if use_raw else column.name
+
+    @classmethod
+    def get_related_field_name(cls, target_field: Type["BaseField"]) -> str:
+        if issubclass(target_field, ormar.fields.ManyToManyField):
+            return cls.get_name()
+        if target_field.virtual:
+            return target_field.related_name or cls.get_name() + "s"
+        return target_field.to.Meta.pkname
+
+    @classmethod
+    def get_filtered_names_to_extract(cls, prefetch_dict: Dict) -> List:
+        related_to_extract = []
+        if prefetch_dict and prefetch_dict is not Ellipsis:
+            related_to_extract = [
+                related
+                for related in cls.extract_related_names()
+                if related in prefetch_dict
+            ]
+        return related_to_extract