Merge pull request #34 from collerek/include_properties_in_dict_and_json

Include properties in dict and json

ormar/__init__.py

@@ -30,7 +30,7 @@ class UndefinedType:  # pragma no cover
 
 Undefined = UndefinedType()
 
-__version__ = "0.4.2"
+__version__ = "0.4.3"
 __all__ = [
     "Integer",
     "BigInteger",

ormar/fields/foreign_key.py

@@ -30,7 +30,9 @@ def create_dummy_model(
     pk_field: Type[Union[BaseField, "ForeignKeyField", "ManyToManyField"]],
 ) -> Type["BaseModel"]:
     fields = {f"{pk_field.name}": (pk_field.__type__, None)}
-    dummy_model = create_model(f"PkOnly{base_model.get_name(lower=False)}", **fields)  # type: ignore
+    dummy_model = create_model(
+        f"PkOnly{base_model.get_name(lower=False)}", **fields  # type: ignore
+    )
     return dummy_model
 
 

ormar/models/newbasemodel.py

@@ -5,6 +5,7 @@ from typing import (
     Any,
     Callable,
     Dict,
+    List,
     Mapping,
     Optional,
     Sequence,
@@ -176,6 +177,24 @@ class NewBaseModel(pydantic.BaseModel, ModelTableProxy, metaclass=ModelMetaclass
     def remove(self, name: "T") -> None:
         self._orm.remove_parent(self, name)
 
+    @classmethod
+    def get_properties(
+        cls,
+        include: Union["AbstractSetIntStr", "MappingIntStrAny"] = None,
+        exclude: Union["AbstractSetIntStr", "MappingIntStrAny"] = None,
+    ) -> List[str]:
+        props = [
+            prop
+            for prop in dir(cls)
+            if isinstance(getattr(cls, prop), property)
+            and prop not in ("__values__", "__fields__", "fields", "pk_column")
+        ]
+        if include:
+            props = [prop for prop in props if prop in include]
+        if exclude:
+            props = [prop for prop in props if prop not in exclude]
+        return props
+
     def dict(  # noqa A003
         self,
         *,
@@ -214,6 +233,12 @@ class NewBaseModel(pydantic.BaseModel, ModelTableProxy, metaclass=ModelMetaclass
                 dict_instance[field] = nested_model.dict(nested=True)
             else:
                 dict_instance[field] = None
+
+        # include model properties as fields
+        props = self.get_properties(include=include, exclude=exclude)
+        if props:
+            dict_instance.update({prop: getattr(self, prop) for prop in props})
+
         return dict_instance
 
     def from_dict(self, value_dict: Dict) -> "NewBaseModel":

tests/test_properties.py

@@ -0,0 +1,66 @@
+import databases
+import pytest
+import sqlalchemy
+
+import ormar
+from tests.settings import DATABASE_URL
+
+database = databases.Database(DATABASE_URL, force_rollback=True)
+metadata = sqlalchemy.MetaData()
+
+
+class Song(ormar.Model):
+    class Meta:
+        tablename = "songs"
+        metadata = metadata
+        database = database
+
+    id: int = ormar.Integer(primary_key=True)
+    name: str = ormar.String(max_length=100)
+    sort_order: int = ormar.Integer()
+
+    @property
+    def sorted_name(self):
+        return f"{self.sort_order}: {self.name}"
+
+    @property
+    def sample(self):
+        return "sample"
+
+    @property
+    def sample2(self):
+        return "sample2"
+
+
+@pytest.fixture(autouse=True, scope="module")
+def create_test_database():
+    engine = sqlalchemy.create_engine(DATABASE_URL)
+    metadata.drop_all(engine)
+    metadata.create_all(engine)
+    yield
+    metadata.drop_all(engine)
+
+
+@pytest.mark.asyncio
+async def test_sort_order_on_main_model():
+    async with database:
+        await Song.objects.create(name="Song 3", sort_order=3)
+        await Song.objects.create(name="Song 1", sort_order=1)
+        await Song.objects.create(name="Song 2", sort_order=2)
+
+        songs = await Song.objects.all()
+        song_dict = [song.dict() for song in songs]
+        assert all('sorted_name' in x for x in song_dict)
+        assert all(x['sorted_name'] == f"{x['sort_order']}: {x['name']}" for x in song_dict)
+        song_json = [song.json() for song in songs]
+        assert all('sorted_name' in x for x in song_json)
+
+        check_include = songs[0].dict(include={"sample"})
+        assert 'sample' in check_include
+        assert 'sample2' not in check_include
+        assert 'sorted_name' not in check_include
+
+        check_include = songs[0].dict(exclude={"sample"})
+        assert 'sample' not in check_include
+        assert 'sample2' in check_include
+        assert 'sorted_name' in check_include