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,
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

View File

@ -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",

View File

@ -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.

View File

@ -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

View File

@ -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.

View File

@ -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)