VimPlug/django-stubs
1
0
Fork 0
mirror of https://github.com/davidhalter/django-stubs.git synced 2025-12-09 05:24:53 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

allow to subclass queryset without loss of typing

Browse Source
This commit is contained in:
Maxim Kurnikov
2019-07-26 22:23:46 +03:00
parent dddcb20fe4
commit 27793ecd32
3 changed files with 42 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
mypy_django_plugin/transformers/querysets.py
View File

@@ -3,6 +3,7 @@ from typing import List, Optional, Sequence, Type, Union, cast
from django.core.exceptions import FieldError
from django.db.models.base import Model
from django.db.models.fields.related import RelatedField
from mypy.newsemanal.typeanal import TypeAnalyser
from mypy.nodes import Expression, NameExpr, TypeInfo
from mypy.plugin import AnalyzeTypeContext, FunctionContext, MethodContext
@@ -10,11 +11,19 @@ from mypy.types import AnyType, Instance
from mypy.types import Type as MypyType
from mypy.types import TypeOfAny
from django.db.models.fields.related import RelatedField
from mypy_django_plugin.django.context import DjangoContext
from mypy_django_plugin.lib import fullnames, helpers
def _extract_model_type_from_queryset(queryset_type: Instance) -> Optional[Instance]:
for base_type in [queryset_type, *queryset_type.type.bases]:
if (len(base_type.args)
and isinstance(base_type.args[0], Instance)
and base_type.args[0].type.has_base(fullnames.MODEL_CLASS_FULLNAME)):
return base_type.args[0]
return None
def determine_proper_manager_type(ctx: FunctionContext) -> MypyType:
default_return_type = ctx.default_return_type
assert isinstance(default_return_type, Instance)
@@ -98,11 +107,10 @@ def extract_proper_type_queryset_values_list(ctx: MethodContext, django_context:
assert isinstance(ctx.type, Instance)
assert isinstance(ctx.default_return_type, Instance)
# bail if queryset of Any or other non-instances
if not isinstance(ctx.type.args[0], Instance):
model_type = _extract_model_type_from_queryset(ctx.type)
if model_type is None:
return AnyType(TypeOfAny.from_omitted_generics)
model_type = ctx.type.args[0]
model_cls = django_context.get_model_class_by_fullname(model_type.type.fullname())
if model_cls is None:
return ctx.default_return_type
@@ -148,11 +156,10 @@ def extract_proper_type_queryset_values(ctx: MethodContext, django_context: Djan
assert isinstance(ctx.type, Instance)
assert isinstance(ctx.default_return_type, Instance)
# if queryset of non-instance type
if not isinstance(ctx.type.args[0], Instance):
model_type = _extract_model_type_from_queryset(ctx.type)
if model_type is None:
return AnyType(TypeOfAny.from_omitted_generics)
model_type = ctx.type.args[0]
model_cls = django_context.get_model_class_by_fullname(model_type.type.fullname())
if model_cls is None:
return ctx.default_return_type