diff --git a/ormar/__init__.py b/ormar/__init__.py index f3cdedb..9bab42f 100644 --- a/ormar/__init__.py +++ b/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", diff --git a/ormar/fields/foreign_key.py b/ormar/fields/foreign_key.py index d2ffb44..f6ad4db 100644 --- a/ormar/fields/foreign_key.py +++ b/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 diff --git a/ormar/models/newbasemodel.py b/ormar/models/newbasemodel.py index 3f48ddf..4cb67fa 100644 --- a/ormar/models/newbasemodel.py +++ b/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": diff --git a/tests/test_properties.py b/tests/test_properties.py new file mode 100644 index 0000000..72398b3 --- /dev/null +++ b/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