Skip to content

Add custom Field processing for mixins used in the Model subclasses #167

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

Merged
merged 1 commit into from
Sep 17, 2019
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 9 additions & 0 deletions mypy_django_plugin/lib/helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,8 @@
)
from mypy.types import AnyType, Instance, NoneTyp, TupleType, Type as MypyType, TypeOfAny, TypedDictType, UnionType

from mypy_django_plugin.lib import fullnames

if TYPE_CHECKING:
from mypy_django_plugin.django.context import DjangoContext

Expand Down Expand Up @@ -247,3 +249,10 @@ def get_typechecker_api(ctx: Union[AttributeContext, MethodContext, FunctionCont
if not isinstance(ctx.api, TypeChecker):
raise ValueError('Not a TypeChecker')
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())
4 changes: 3 additions & 1 deletion mypy_django_plugin/transformers/fields.py
Original file line number Diff line number Diff line change
Expand Up @@ -116,7 +116,9 @@ def transform_into_proper_return_type(ctx: FunctionContext, django_context: Djan
assert isinstance(default_return_type, Instance)

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
return ctx.default_return_type
assert isinstance(outer_model_info, TypeInfo)
Expand Down
15 changes: 14 additions & 1 deletion mypy_django_plugin/transformers/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -212,6 +212,18 @@ def run_with_model_cls(self, model_cls: Type[Model]) -> None:
]))


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,
django_context: DjangoContext) -> None:
initializers = [
Expand All @@ -220,7 +232,8 @@ def process_model_class(ctx: ClassDefContext,
AddRelatedModelsId,
AddManagers,
AddExtraFieldMethods,
AddMetaOptionsAttribute
AddMetaOptionsAttribute,
RecordAllModelMixins,
]
for initializer_cls in initializers:
try:
Expand Down
16 changes: 16 additions & 0 deletions test-data/typecheck/fields/test_base.yml
Original file line number Diff line number Diff line change
Expand Up @@ -132,3 +132,19 @@
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]'

- 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