some types improvements

This commit is contained in:
collerek
2021-03-19 18:40:46 +01:00
parent 9c091afe35
commit 859ed5d1fc
6 changed files with 45 additions and 42 deletions

View File

@ -14,7 +14,7 @@ from typing import (
TYPE_CHECKING, TYPE_CHECKING,
Type, Type,
Union, Union,
cast, cast, no_type_check,
) )
try: try:
@ -47,6 +47,7 @@ from ormar.relations.relation_manager import RelationsManager
if TYPE_CHECKING: # pragma no cover if TYPE_CHECKING: # pragma no cover
from ormar.models import Model from ormar.models import Model
from ormar.signals import SignalEmitter from ormar.signals import SignalEmitter
from ormar.queryset import QuerySet
IntStr = Union[int, str] IntStr = Union[int, str]
DictStrAny = Dict[str, Any] DictStrAny = Dict[str, Any]
@ -67,6 +68,7 @@ class NewBaseModel(pydantic.BaseModel, ModelTableProxy, metaclass=ModelMetaclass
__slots__ = ("_orm_id", "_orm_saved", "_orm", "_pk_column") __slots__ = ("_orm_id", "_orm_saved", "_orm", "_pk_column")
if TYPE_CHECKING: # pragma no cover if TYPE_CHECKING: # pragma no cover
pk: Any
__model_fields__: Dict[str, BaseField] __model_fields__: Dict[str, BaseField]
__table__: sqlalchemy.Table __table__: sqlalchemy.Table
__fields__: Dict[str, pydantic.fields.ModelField] __fields__: Dict[str, pydantic.fields.ModelField]
@ -77,6 +79,7 @@ class NewBaseModel(pydantic.BaseModel, ModelTableProxy, metaclass=ModelMetaclass
__database__: databases.Database __database__: databases.Database
_orm_relationship_manager: AliasManager _orm_relationship_manager: AliasManager
_orm: RelationsManager _orm: RelationsManager
_orm_id: int
_orm_saved: bool _orm_saved: bool
_related_names: Optional[Set] _related_names: Optional[Set]
_related_names_hash: str _related_names_hash: str
@ -229,7 +232,7 @@ class NewBaseModel(pydantic.BaseModel, ModelTableProxy, metaclass=ModelMetaclass
super().__setattr__(name, value) super().__setattr__(name, value)
self.set_save_status(False) 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 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 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": if item == "pk":
return object.__getattribute__(self, "__dict__").get(self.Meta.pkname, None) return object.__getattribute__(self, "__dict__").get(self.Meta.pkname, None)
if item in object.__getattribute__(self, "extract_related_names")(): if item in object.__getattribute__(self, "extract_related_names")():
return object.__getattribute__( return self._extract_related_model_instead_of_field(item)
self, "_extract_related_model_instead_of_field"
)(item)
if item in object.__getattribute__(self, "extract_through_names")(): if item in object.__getattribute__(self, "extract_through_names")():
return object.__getattribute__( return self._extract_related_model_instead_of_field(item)
self, "_extract_related_model_instead_of_field"
)(item)
if item in object.__getattribute__(self, "Meta").property_fields: if item in object.__getattribute__(self, "Meta").property_fields:
value = object.__getattribute__(self, item) value = object.__getattribute__(self, item)
return value() if callable(value) else value return value() if callable(value) else value

View File

@ -23,6 +23,7 @@ quick_access_set = {
"_extract_nested_models", "_extract_nested_models",
"_extract_nested_models_from_list", "_extract_nested_models_from_list",
"_extract_own_model_fields", "_extract_own_model_fields",
"_extract_related_model_instead_of_field",
"_get_related_not_excluded_fields", "_get_related_not_excluded_fields",
"_get_value", "_get_value",
"_is_conversion_to_json_needed", "_is_conversion_to_json_needed",

View File

@ -302,7 +302,7 @@ class QuerysetProxy(Generic[T]):
self._register_related(get) self._register_related(get)
return 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. Returns all rows from a database for given model for set filter options.

View File

@ -1,5 +1,6 @@
from enum import Enum 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 import ormar # noqa I100
from ormar.exceptions import RelationshipInstanceError # 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 if TYPE_CHECKING: # pragma no cover
from ormar.relations import RelationsManager 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): class RelationType(Enum):
@ -25,7 +29,7 @@ class RelationType(Enum):
THROUGH = 4 THROUGH = 4
class Relation: class Relation(Generic[T]):
""" """
Keeps related Models and handles adding/removing of the children. Keeps related Models and handles adding/removing of the children.
""" """
@ -35,7 +39,7 @@ class Relation:
manager: "RelationsManager", manager: "RelationsManager",
type_: RelationType, type_: RelationType,
field_name: str, field_name: str,
to: Type["Model"], to: Type["T"],
through: Type["Model"] = None, through: Type["Model"] = None,
) -> None: ) -> None:
""" """
@ -59,7 +63,7 @@ class Relation:
self._owner: "Model" = manager.owner self._owner: "Model" = manager.owner
self._type: RelationType = type_ self._type: RelationType = type_
self._to_remove: Set = set() self._to_remove: Set = set()
self.to: Type["Model"] = to self.to: Type["T"] = to
self._through = through self._through = through
self.field_name: str = field_name self.field_name: str = field_name
self.related_models: Optional[Union[RelationProxy, "Model"]] = ( self.related_models: Optional[Union[RelationProxy, "Model"]] = (
@ -73,7 +77,8 @@ class Relation:
self.related_models = None self.related_models = None
self._owner.__dict__[self.field_name] = None self._owner.__dict__[self.field_name] = None
elif self.related_models is not 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 self._owner.__dict__[self.field_name] = None
@property @property

View File

@ -27,11 +27,11 @@ class RelationProxy(Generic[T], list):
data_: Any = None, data_: Any = None,
) -> None: ) -> None:
super().__init__(data_ or ()) super().__init__(data_ or ())
self.relation: "Relation" = relation self.relation: "Relation[T]" = relation
self.type_: "RelationType" = type_ self.type_: "RelationType" = type_
self.field_name = field_name self.field_name = field_name
self._owner: "Model" = self.relation.manager.owner 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_ relation=self.relation, to=to, type_=type_
) )
self._related_field_name: Optional[str] = None self._related_field_name: Optional[str] = None
@ -70,7 +70,7 @@ class RelationProxy(Generic[T], list):
return getattr(self.queryset_proxy, item) return getattr(self.queryset_proxy, item)
return super().__getattribute__(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. Delegates calls for non existing attributes to QuerySetProxy.
@ -114,7 +114,7 @@ class RelationProxy(Generic[T], list):
"You cannot query relationships from unsaved model." "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 Creates new QuerySet with relation model and pre filters it with currents
parent model primary key, so all queries by definition are already related parent model primary key, so all queries by definition are already related
@ -138,7 +138,7 @@ class RelationProxy(Generic[T], list):
return queryset return queryset
async def remove( # type: ignore async def remove( # type: ignore
self, item: "Model", keep_reversed: bool = True self, item: "T", keep_reversed: bool = True
) -> None: ) -> None:
""" """
Removes the related from relation with parent. Removes the related from relation with parent.
@ -189,7 +189,7 @@ class RelationProxy(Generic[T], list):
relation_name=self.field_name, 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. Adds child model to relation.

View File

@ -73,31 +73,29 @@ async def test_types() -> None:
authors = publisher2.authors authors = publisher2.authors
assert authors[0] == author assert authors[0] == author
for author in authors: for author in authors:
if TYPE_CHECKING: # pragma: no cover pass
reveal_type(author) # iter of relation proxy # if TYPE_CHECKING: # pragma: no cover
# reveal_type(author) # iter of relation proxy
book = await Book.objects.create(title="Test", author=author) book = await Book.objects.create(title="Test", author=author)
book2 = await Book.objects.select_related("author").get() book2 = await Book.objects.select_related("author").get()
books = await Book.objects.select_related("author").all() books = await Book.objects.select_related("author").all()
author_books = await author.books.all() author_books = await author.books.all()
assert book.author.name == "Test Author" assert book.author.name == "Test Author"
assert book2.author.name == "Test Author" assert book2.author.name == "Test Author"
if TYPE_CHECKING: # pragma: no cover # if TYPE_CHECKING: # pragma: no cover
reveal_type(publisher) # model method # reveal_type(publisher) # model method
reveal_type(publishers) # many to many # reveal_type(publishers) # many to many
reveal_type(publishers[0]) # item in m2m list # reveal_type(publishers[0]) # item in m2m list
# getting relation without __getattribute__ # # getting relation without __getattribute__
reveal_type(authors) # reverse many to many # TODO: wrong # reveal_type(authors) # reverse many to many # TODO: wrong
reveal_type(book2) # queryset get # reveal_type(book2) # queryset get
reveal_type(books) # queryset all # reveal_type(books) # queryset all
reveal_type(book) # queryset - create # reveal_type(book) # queryset - create
reveal_type(query) # queryset itself # reveal_type(query) # queryset itself
reveal_type(book.author) # fk # reveal_type(book.author) # fk
reveal_type( # reveal_type(author.books) # reverse fk relation proxy # TODO: wrong
author.books.queryset_proxy # reveal_type(author) # another test for queryset get different model
) # queryset in querysetproxy # TODO: wrong # reveal_type(book.author.name) # field on related model
reveal_type(author.books) # reverse fk # TODO: wrong # reveal_type(author_books) # querysetproxy result for fk # TODO: wrong
reveal_type(author) # another test for queryset get different model # reveal_type(author_books[0]) # item in qs proxy for fk # TODO: wrong
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
assert_type(book) assert_type(book)