From 859ed5d1fca42679f2f4294c0b576fc1c3a7c293 Mon Sep 17 00:00:00 2001 From: collerek Date: Fri, 19 Mar 2021 18:40:46 +0100 Subject: [PATCH] some types improvements --- ormar/models/newbasemodel.py | 15 ++++++----- ormar/models/quick_access_views.py | 1 + ormar/relations/querysetproxy.py | 2 +- ormar/relations/relation.py | 17 ++++++++----- ormar/relations/relation_proxy.py | 12 ++++----- tests/test_types.py | 40 ++++++++++++++---------------- 6 files changed, 45 insertions(+), 42 deletions(-) diff --git a/ormar/models/newbasemodel.py b/ormar/models/newbasemodel.py index 02a2663..0943457 100644 --- a/ormar/models/newbasemodel.py +++ b/ormar/models/newbasemodel.py @@ -14,7 +14,7 @@ from typing import ( TYPE_CHECKING, Type, Union, - cast, + cast, no_type_check, ) try: @@ -47,6 +47,7 @@ from ormar.relations.relation_manager import RelationsManager if TYPE_CHECKING: # pragma no cover from ormar.models import Model from ormar.signals import SignalEmitter + from ormar.queryset import QuerySet IntStr = Union[int, str] DictStrAny = Dict[str, Any] @@ -67,6 +68,7 @@ class NewBaseModel(pydantic.BaseModel, ModelTableProxy, metaclass=ModelMetaclass __slots__ = ("_orm_id", "_orm_saved", "_orm", "_pk_column") if TYPE_CHECKING: # pragma no cover + pk: Any __model_fields__: Dict[str, BaseField] __table__: sqlalchemy.Table __fields__: Dict[str, pydantic.fields.ModelField] @@ -77,6 +79,7 @@ class NewBaseModel(pydantic.BaseModel, ModelTableProxy, metaclass=ModelMetaclass __database__: databases.Database _orm_relationship_manager: AliasManager _orm: RelationsManager + _orm_id: int _orm_saved: bool _related_names: Optional[Set] _related_names_hash: str @@ -229,7 +232,7 @@ class NewBaseModel(pydantic.BaseModel, ModelTableProxy, metaclass=ModelMetaclass super().__setattr__(name, value) self.set_save_status(False) - def __getattribute__(self, item: str) -> Any: # noqa: CCR001 + def __getattribute__(self, item: str): # noqa: CCR001 """ Because we need to overwrite getting the attribute by ormar instead of pydantic as well as returning related models and not the value stored on the model the @@ -265,13 +268,9 @@ class NewBaseModel(pydantic.BaseModel, ModelTableProxy, metaclass=ModelMetaclass if item == "pk": return object.__getattribute__(self, "__dict__").get(self.Meta.pkname, None) if item in object.__getattribute__(self, "extract_related_names")(): - return object.__getattribute__( - self, "_extract_related_model_instead_of_field" - )(item) + return self._extract_related_model_instead_of_field(item) if item in object.__getattribute__(self, "extract_through_names")(): - return object.__getattribute__( - self, "_extract_related_model_instead_of_field" - )(item) + return self._extract_related_model_instead_of_field(item) if item in object.__getattribute__(self, "Meta").property_fields: value = object.__getattribute__(self, item) return value() if callable(value) else value diff --git a/ormar/models/quick_access_views.py b/ormar/models/quick_access_views.py index c960898..dd4cb5a 100644 --- a/ormar/models/quick_access_views.py +++ b/ormar/models/quick_access_views.py @@ -23,6 +23,7 @@ quick_access_set = { "_extract_nested_models", "_extract_nested_models_from_list", "_extract_own_model_fields", + "_extract_related_model_instead_of_field", "_get_related_not_excluded_fields", "_get_value", "_is_conversion_to_json_needed", diff --git a/ormar/relations/querysetproxy.py b/ormar/relations/querysetproxy.py index a97c52a..7ac1c3e 100644 --- a/ormar/relations/querysetproxy.py +++ b/ormar/relations/querysetproxy.py @@ -302,7 +302,7 @@ class QuerysetProxy(Generic[T]): self._register_related(get) return get - async def all(self, **kwargs: Any) -> Sequence[Optional["T"]]: # noqa: A003 + async def all(self, **kwargs: Any) -> List[Optional["T"]]: # noqa: A003 """ Returns all rows from a database for given model for set filter options. diff --git a/ormar/relations/relation.py b/ormar/relations/relation.py index 3d30423..96031ec 100644 --- a/ormar/relations/relation.py +++ b/ormar/relations/relation.py @@ -1,5 +1,6 @@ from enum import Enum -from typing import List, Optional, Set, TYPE_CHECKING, Type, Union +from typing import Generic, List, Optional, Set, TYPE_CHECKING, Type, TypeVar, Union, \ + cast import ormar # noqa I100 from ormar.exceptions import RelationshipInstanceError # noqa I100 @@ -7,7 +8,10 @@ from ormar.relations.relation_proxy import RelationProxy if TYPE_CHECKING: # pragma no cover from ormar.relations import RelationsManager - from ormar.models import Model, NewBaseModel + from ormar.models import Model, NewBaseModel, T + from ormar.relations.relation_proxy import RelationProxy +else: + T = TypeVar("T", bound="Model") class RelationType(Enum): @@ -25,7 +29,7 @@ class RelationType(Enum): THROUGH = 4 -class Relation: +class Relation(Generic[T]): """ Keeps related Models and handles adding/removing of the children. """ @@ -35,7 +39,7 @@ class Relation: manager: "RelationsManager", type_: RelationType, field_name: str, - to: Type["Model"], + to: Type["T"], through: Type["Model"] = None, ) -> None: """ @@ -59,7 +63,7 @@ class Relation: self._owner: "Model" = manager.owner self._type: RelationType = type_ self._to_remove: Set = set() - self.to: Type["Model"] = to + self.to: Type["T"] = to self._through = through self.field_name: str = field_name self.related_models: Optional[Union[RelationProxy, "Model"]] = ( @@ -73,7 +77,8 @@ class Relation: self.related_models = None self._owner.__dict__[self.field_name] = None elif self.related_models is not None: - self.related_models._clear() + related_models = cast("RelationProxy", self.related_models) + related_models._clear() self._owner.__dict__[self.field_name] = None @property diff --git a/ormar/relations/relation_proxy.py b/ormar/relations/relation_proxy.py index c2b39a1..55a8225 100644 --- a/ormar/relations/relation_proxy.py +++ b/ormar/relations/relation_proxy.py @@ -27,11 +27,11 @@ class RelationProxy(Generic[T], list): data_: Any = None, ) -> None: super().__init__(data_ or ()) - self.relation: "Relation" = relation + self.relation: "Relation[T]" = relation self.type_: "RelationType" = type_ self.field_name = field_name self._owner: "Model" = self.relation.manager.owner - self.queryset_proxy: QuerysetProxy = QuerysetProxy( + self.queryset_proxy: QuerysetProxy[T] = QuerysetProxy[T]( relation=self.relation, to=to, type_=type_ ) self._related_field_name: Optional[str] = None @@ -70,7 +70,7 @@ class RelationProxy(Generic[T], list): return getattr(self.queryset_proxy, item) return super().__getattribute__(item) - def __getattr__(self, item: str) -> "T": + def __getattr__(self, item: str) -> Any: """ Delegates calls for non existing attributes to QuerySetProxy. @@ -114,7 +114,7 @@ class RelationProxy(Generic[T], list): "You cannot query relationships from unsaved model." ) - def _set_queryset(self) -> "QuerySet": + def _set_queryset(self) -> "QuerySet[T]": """ Creates new QuerySet with relation model and pre filters it with currents parent model primary key, so all queries by definition are already related @@ -138,7 +138,7 @@ class RelationProxy(Generic[T], list): return queryset async def remove( # type: ignore - self, item: "Model", keep_reversed: bool = True + self, item: "T", keep_reversed: bool = True ) -> None: """ Removes the related from relation with parent. @@ -189,7 +189,7 @@ class RelationProxy(Generic[T], list): relation_name=self.field_name, ) - async def add(self, item: "Model", **kwargs: Any) -> None: + async def add(self, item: "T", **kwargs: Any) -> None: """ Adds child model to relation. diff --git a/tests/test_types.py b/tests/test_types.py index 29adfc1..ba819ed 100644 --- a/tests/test_types.py +++ b/tests/test_types.py @@ -73,31 +73,29 @@ async def test_types() -> None: authors = publisher2.authors assert authors[0] == author for author in authors: - if TYPE_CHECKING: # pragma: no cover - reveal_type(author) # iter of relation proxy + pass + # if TYPE_CHECKING: # pragma: no cover + # reveal_type(author) # iter of relation proxy book = await Book.objects.create(title="Test", author=author) book2 = await Book.objects.select_related("author").get() books = await Book.objects.select_related("author").all() author_books = await author.books.all() assert book.author.name == "Test Author" assert book2.author.name == "Test Author" - if TYPE_CHECKING: # pragma: no cover - reveal_type(publisher) # model method - reveal_type(publishers) # many to many - reveal_type(publishers[0]) # item in m2m list - # getting relation without __getattribute__ - reveal_type(authors) # reverse many to many # TODO: wrong - reveal_type(book2) # queryset get - reveal_type(books) # queryset all - reveal_type(book) # queryset - create - reveal_type(query) # queryset itself - reveal_type(book.author) # fk - reveal_type( - author.books.queryset_proxy - ) # queryset in querysetproxy # TODO: wrong - reveal_type(author.books) # reverse fk # TODO: wrong - reveal_type(author) # another test for queryset get different model - reveal_type(book.author.name) # field on related model - reveal_type(author_books) # querysetproxy for fk # TODO: wrong - reveal_type(author_books[0]) # item i qs proxy for fk # TODO: wrong + # if TYPE_CHECKING: # pragma: no cover + # reveal_type(publisher) # model method + # reveal_type(publishers) # many to many + # reveal_type(publishers[0]) # item in m2m list + # # getting relation without __getattribute__ + # reveal_type(authors) # reverse many to many # TODO: wrong + # reveal_type(book2) # queryset get + # reveal_type(books) # queryset all + # reveal_type(book) # queryset - create + # reveal_type(query) # queryset itself + # reveal_type(book.author) # fk + # reveal_type(author.books) # reverse fk relation proxy # TODO: wrong + # reveal_type(author) # another test for queryset get different model + # reveal_type(book.author.name) # field on related model + # reveal_type(author_books) # querysetproxy result for fk # TODO: wrong + # reveal_type(author_books[0]) # item in qs proxy for fk # TODO: wrong assert_type(book)