allow passing a dict and set to fields and exclude_fields, store it as dict

ormar/models/modelproxy.py

@@ -1,6 +1,16 @@
 import inspect
 from collections import OrderedDict
-from typing import Dict, List, Sequence, Set, TYPE_CHECKING, Type, TypeVar, Union
+from typing import (
+    Dict,
+    List,
+    Optional,
+    Sequence,
+    Set,
+    TYPE_CHECKING,
+    Type,
+    TypeVar,
+    Union,
+)
 
 import ormar
 from ormar.exceptions import RelationshipInstanceError
@@ -65,14 +75,14 @@ class ModelTableProxy:
     @classmethod
     def get_column_alias(cls, field_name: str) -> str:
         field = cls.Meta.model_fields.get(field_name)
-        if field and field.alias is not None:
+        if field is not None and field.alias is not None:
             return field.alias
         return 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 and field.alias == alias:
+            if field is not None and field.alias == alias:
                 return field_name
         return alias  # if not found it's not an alias but actual name
 
@@ -211,59 +221,13 @@ class ModelTableProxy:
                 )
         return other
 
-    @staticmethod
-    def _get_not_nested_columns_from_fields(
-        model: Type["Model"],
-        fields: List,
-        exclude_fields: List,
-        column_names: List[str],
-        use_alias: bool = False,
-    ) -> List[str]:
-        fields = [model.get_column_alias(k) if not use_alias else k for k in fields]
-        fields = fields or column_names
-        exclude_fields = [
-            model.get_column_alias(k) if not use_alias else k for k in exclude_fields
-        ]
-        columns = [
-            name
-            for name in fields
-            if "__" not in name and name in column_names and name not in exclude_fields
-        ]
-        return columns
-
-    @staticmethod
-    def _get_nested_columns_from_fields(
-        model: Type["Model"],
-        fields: List,
-        exclude_fields: List,
-        column_names: List[str],
-        use_alias: bool = False,
-    ) -> List[str]:
-        model_name = f"{model.get_name()}__"
-        columns = [
-            name[(name.find(model_name) + len(model_name)) :]  # noqa: E203
-            for name in fields
-            if f"{model.get_name()}__" in name
-        ]
-        columns = columns or column_names
-        exclude_columns = [
-            name[(name.find(model_name) + len(model_name)) :]  # noqa: E203
-            for name in exclude_fields
-            if f"{model.get_name()}__" in name
-        ]
-        columns = [model.get_column_alias(k) if not use_alias else k for k in columns]
-        exclude_columns = [
-            model.get_column_alias(k) if not use_alias else k for k in exclude_columns
-        ]
-        return [column for column in columns if column not in exclude_columns]
-
     @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 not use_alias
+            if use_alias
             else model.Meta.pkname
         )
         if pk_alias not in columns:
@@ -273,34 +237,30 @@ class ModelTableProxy:
     @staticmethod
     def own_table_columns(
         model: Type["Model"],
-        fields: List,
-        exclude_fields: List,
-        nested: bool = False,
+        fields: Optional[Union[Set, Dict]],
+        exclude_fields: Optional[Union[Set, Dict]],
         use_alias: bool = False,
     ) -> List[str]:
-        column_names = [
-            model.get_column_name_from_alias(col.name) if use_alias else col.name
+        columns = [
+            model.get_column_name_from_alias(col.name) if not use_alias else col.name
             for col in model.Meta.table.columns
         ]
-        if not fields and not exclude_fields:
-            return column_names
-
-        if not nested:
-            columns = ModelTableProxy._get_not_nested_columns_from_fields(
-                model=model,
-                fields=fields,
-                exclude_fields=exclude_fields,
-                column_names=column_names,
-                use_alias=use_alias,
-            )
-        else:
-            columns = ModelTableProxy._get_nested_columns_from_fields(
-                model=model,
-                fields=fields,
-                exclude_fields=exclude_fields,
-                column_names=column_names,
-                use_alias=use_alias,
-            )
+        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
         columns = ModelTableProxy._populate_pk_column(