add custom Field processing for mixins used in the Model subclasses (#167)

This commit is contained in:
Maxim Kurnikov
2019-09-18 01:40:41 +03:00
committed by GitHub
parent 64720f492c
commit 23ad65033b
4 changed files with 42 additions and 2 deletions

View File

@@ -11,6 +11,8 @@ from mypy.plugin import (
) )
from mypy.types import AnyType, Instance, NoneTyp, TupleType, Type as MypyType, TypeOfAny, TypedDictType, UnionType from mypy.types import AnyType, Instance, NoneTyp, TupleType, Type as MypyType, TypeOfAny, TypedDictType, UnionType
from mypy_django_plugin.lib import fullnames
if TYPE_CHECKING: if TYPE_CHECKING:
from mypy_django_plugin.django.context import DjangoContext from mypy_django_plugin.django.context import DjangoContext
@@ -247,3 +249,10 @@ def get_typechecker_api(ctx: Union[AttributeContext, MethodContext, FunctionCont
if not isinstance(ctx.api, TypeChecker): if not isinstance(ctx.api, TypeChecker):
raise ValueError('Not a TypeChecker') raise ValueError('Not a TypeChecker')
return cast(TypeChecker, ctx.api) return cast(TypeChecker, ctx.api)
def get_all_model_mixins(api: TypeChecker) -> Set[str]:
basemodel_info = lookup_fully_qualified_typeinfo(api, fullnames.MODEL_CLASS_FULLNAME)
if basemodel_info is None:
return set()
return set(get_django_metadata(basemodel_info).get('model_mixins', dict).keys())

View File

@@ -116,7 +116,9 @@ def transform_into_proper_return_type(ctx: FunctionContext, django_context: Djan
assert isinstance(default_return_type, Instance) assert isinstance(default_return_type, Instance)
outer_model_info = helpers.get_typechecker_api(ctx).scope.active_class() outer_model_info = helpers.get_typechecker_api(ctx).scope.active_class()
if not outer_model_info or not outer_model_info.has_base(fullnames.MODEL_CLASS_FULLNAME): if (outer_model_info is None
or not outer_model_info.has_base(fullnames.MODEL_CLASS_FULLNAME)
and outer_model_info.fullname() not in helpers.get_all_model_mixins(helpers.get_typechecker_api(ctx))):
# not inside models.Model class # not inside models.Model class
return ctx.default_return_type return ctx.default_return_type
assert isinstance(outer_model_info, TypeInfo) assert isinstance(outer_model_info, TypeInfo)

View File

@@ -212,6 +212,18 @@ class AddMetaOptionsAttribute(ModelClassInitializer):
])) ]))
class RecordAllModelMixins(ModelClassInitializer):
def run(self) -> None:
basemodel_info = self.lookup_typeinfo_or_incomplete_defn_error(fullnames.MODEL_CLASS_FULLNAME)
basemodel_metadata = helpers.get_django_metadata(basemodel_info)
if 'model_mixins' not in basemodel_metadata:
basemodel_metadata['model_mixins'] = {}
for base_info in self.model_classdef.info.mro[1:]:
if base_info.fullname() != 'builtins.object':
basemodel_metadata['model_mixins'][base_info.fullname()] = 1
def process_model_class(ctx: ClassDefContext, def process_model_class(ctx: ClassDefContext,
django_context: DjangoContext) -> None: django_context: DjangoContext) -> None:
initializers = [ initializers = [
@@ -220,7 +232,8 @@ def process_model_class(ctx: ClassDefContext,
AddRelatedModelsId, AddRelatedModelsId,
AddManagers, AddManagers,
AddExtraFieldMethods, AddExtraFieldMethods,
AddMetaOptionsAttribute AddMetaOptionsAttribute,
RecordAllModelMixins,
] ]
for initializer_cls in initializers: for initializer_cls in initializers:
try: try:

View File

@@ -132,3 +132,19 @@
myfield: models.IntegerField[int, int] myfield: models.IntegerField[int, int]
reveal_type(MyClass.myfield) # N: Revealed type is 'django.db.models.fields.IntegerField[builtins.int, builtins.int]' reveal_type(MyClass.myfield) # N: Revealed type is 'django.db.models.fields.IntegerField[builtins.int, builtins.int]'
reveal_type(MyClass().myfield) # N: Revealed type is 'django.db.models.fields.IntegerField[builtins.int, builtins.int]' reveal_type(MyClass().myfield) # N: Revealed type is 'django.db.models.fields.IntegerField[builtins.int, builtins.int]'
- case: fields_inside_mixins_used_in_model_subclasses_resolved_as_primitives
main: |
from myapp.models import MyModel, AuthMixin
reveal_type(MyModel().username) # N: Revealed type is 'builtins.str*'
installed_apps:
- myapp
files:
- path: myapp/__init__.py
- path: myapp/models.py
content: |
from django.db import models
class AuthMixin:
username = models.CharField(max_length=100)
class MyModel(AuthMixin, models.Model):
pass