From 9838547c4ff1a4b0989892f2c0924fdbbfbff6cc Mon Sep 17 00:00:00 2001
From: collerek <collerek@gmail.com>
Date: Sun, 6 Dec 2020 08:23:57 +0100
Subject: [PATCH] some cleanup and refactoring

---
 ormar/models/newbasemodel.py              | 29 +++++++++++------------
 tests/test_excluding_fields_in_fastapi.py | 10 +++++---
 2 files changed, 21 insertions(+), 18 deletions(-)

diff --git a/ormar/models/newbasemodel.py b/ormar/models/newbasemodel.py
index 9105f7b..45e0fa0 100644
--- a/ormar/models/newbasemodel.py
+++ b/ormar/models/newbasemodel.py
@@ -47,11 +47,7 @@ if TYPE_CHECKING:  # pragma no cover
 class NewBaseModel(
     pydantic.BaseModel, ModelTableProxy, Excludable, metaclass=ModelMetaclass
 ):
-    __slots__ = (
-        "_orm_id",
-        "_orm_saved",
-        "_orm",
-    )
+    __slots__ = ("_orm_id", "_orm_saved", "_orm", "_pk_column")
 
     if TYPE_CHECKING:  # pragma no cover
         __model_fields__: Dict[str, Type[BaseField]]
@@ -75,6 +71,7 @@ class NewBaseModel(
     def __init__(self, *args: Any, **kwargs: Any) -> None:  # type: ignore
         object.__setattr__(self, "_orm_id", uuid.uuid4().hex)
         object.__setattr__(self, "_orm_saved", False)
+        object.__setattr__(self, "_pk_column", None)
         object.__setattr__(
             self,
             "_orm",
@@ -94,13 +91,8 @@ class NewBaseModel(
         if "pk" in kwargs:
             kwargs[self.Meta.pkname] = kwargs.pop("pk")
 
-        # remove property fields values from validation
-        kwargs = {
-            k: v
-            for k, v in kwargs.items()
-            if k not in object.__getattribute__(self, "Meta").property_fields
-        }
         # build the models to set them and validate but don't register
+        # also remove property fields values from validation
         try:
             new_kwargs: Dict[str, Any] = {
                 k: self._convert_json(
@@ -111,14 +103,15 @@ class NewBaseModel(
                     "dumps",
                 )
                 for k, v in kwargs.items()
+                if k not in object.__getattribute__(self, "Meta").property_fields
             }
         except KeyError as e:
             raise ModelError(
                 f"Unknown field '{e.args[0]}' for model {self.get_name(lower=False)}"
             )
 
-        # explicitly set None to excluded fields with default
-        # as pydantic populates them with default
+        # explicitly set None to excluded fields
+        # as pydantic populates them with default if set
         for field_to_nullify in excluded:
             new_kwargs[field_to_nullify] = None
 
@@ -195,7 +188,8 @@ class NewBaseModel(
         return (
             self._orm_id == other._orm_id
             or (self.pk == other.pk and self.pk is not None)
-            or self.dict() == other.dict()
+            or self.dict(exclude=self.extract_related_names())
+            == other.dict(exclude=other.extract_related_names())
         )
 
     @classmethod
@@ -207,7 +201,12 @@ class NewBaseModel(
 
     @property
     def pk_column(self) -> sqlalchemy.Column:
-        return self.Meta.table.primary_key.columns.values()[0]
+        if object.__getattribute__(self, "_pk_column") is not None:
+            return object.__getattribute__(self, "_pk_column")
+        pk_columns = self.Meta.table.primary_key.columns.values()
+        pk_col = pk_columns[0]
+        object.__setattr__(self, "_pk_column", pk_col)
+        return pk_col
 
     @property
     def saved(self) -> bool:
diff --git a/tests/test_excluding_fields_in_fastapi.py b/tests/test_excluding_fields_in_fastapi.py
index a75d0e1..3e3349b 100644
--- a/tests/test_excluding_fields_in_fastapi.py
+++ b/tests/test_excluding_fields_in_fastapi.py
@@ -219,7 +219,7 @@ def test_excluding_fields_in_endpoints():
         assert isinstance(user_instance.timestamp, datetime.datetime)
         assert user_instance.timestamp == timestamp
 
-        response = client.post("/users4/", json=user)
+        response = client.post("/users4/", json=user3)
         assert list(response.json().keys()) == [
             "id",
             "email",
@@ -228,8 +228,12 @@ def test_excluding_fields_in_endpoints():
             "category",
             "timestamp",
         ]
-        assert response.json().get("timestamp") != str(timestamp).replace(" ", "T")
-        assert response.json().get("timestamp") is not None
+        assert (
+            datetime.datetime.strptime(
+                response.json().get("timestamp"), "%Y-%m-%dT%H:%M:%S.%f"
+            )
+            == timestamp
+        )
 
 
 def test_adding_fields_in_endpoints():