Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Added support to create Generic models #987

Open
wants to merge 8 commits into
base: develop
Choose a base branch
from
10 changes: 7 additions & 3 deletions tortoise/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
Type,
TypeVar,
Union,
Generic,
)

from pypika import Order, Query, Table
Expand Down Expand Up @@ -630,9 +631,6 @@ def __search_for_field_attributes(base: Type, attrs: dict) -> None:
meta.finalise_fields()
return new_class

def __getitem__(cls: Type[MODEL], key: Any) -> QuerySetSingle[MODEL]: # type: ignore
return cls._getbypk(key) # type: ignore


class Model(metaclass=ModelMeta):
"""
Expand Down Expand Up @@ -663,6 +661,12 @@ def __init__(self, **kwargs: Any) -> None:
else:
setattr(self, key, field_object.default)

def __class_getitem__(cls: Type[MODEL], key: Any) -> QuerySetSingle[MODEL]: # type: ignore
if ((isinstance(key, tuple) and inspect.isclass(key[0])) or inspect.isclass(key)) and issubclass(cls, Generic):
return Generic.__class_getitem__.__func__(cls, key)
return cls._getbypk(key)


def _set_kwargs(self, kwargs: dict) -> Set[str]:
meta = self._meta

Expand Down