diff --git a/django-stubs/http/request.pyi b/django-stubs/http/request.pyi index 035727a..eae05c3 100644 --- a/django-stubs/http/request.pyi +++ b/django-stubs/http/request.pyi @@ -17,8 +17,10 @@ from typing import ( ) from django.contrib.sessions.backends.base import SessionBase +from django.db.models.base import Model from django.utils.datastructures import CaseInsensitiveMapping, ImmutableList, MultiValueDict +from django.contrib.auth.base_user import AbstractBaseUser from django.core.files import uploadedfile, uploadhandler from django.urls import ResolverMatch @@ -49,6 +51,7 @@ class HttpRequest(BytesIO): resolver_match: ResolverMatch = ... content_type: Optional[str] = ... content_params: Optional[Dict[str, str]] = ... + user: AbstractBaseUser session: SessionBase encoding: Optional[str] = ... upload_handlers: UploadHandlerList = ... diff --git a/mypy_django_plugin/lib/fullnames.py b/mypy_django_plugin/lib/fullnames.py index 87d9b43..15e5ce8 100644 --- a/mypy_django_plugin/lib/fullnames.py +++ b/mypy_django_plugin/lib/fullnames.py @@ -36,3 +36,4 @@ RELATED_FIELDS_CLASSES = { MIGRATION_CLASS_FULLNAME = 'django.db.migrations.migration.Migration' OPTIONS_CLASS_FULLNAME = 'django.db.models.options.Options' +HTTPREQUEST_CLASS_FULLNAME = 'django.http.request.HttpRequest' diff --git a/mypy_django_plugin/lib/helpers.py b/mypy_django_plugin/lib/helpers.py index 39c8bc4..b0277fc 100644 --- a/mypy_django_plugin/lib/helpers.py +++ b/mypy_django_plugin/lib/helpers.py @@ -49,7 +49,7 @@ def lookup_fully_qualified_typeinfo(api: TypeChecker, fullname: str) -> Optional return node -def lookup_class_typeinfo(api: TypeChecker, klass: type) -> TypeInfo: +def lookup_class_typeinfo(api: TypeChecker, klass: type) -> Optional[TypeInfo]: fullname = get_class_fullname(klass) field_info = lookup_fully_qualified_typeinfo(api, fullname) return field_info diff --git a/mypy_django_plugin/main.py b/mypy_django_plugin/main.py index 9a20873..254f642 100644 --- a/mypy_django_plugin/main.py +++ b/mypy_django_plugin/main.py @@ -11,7 +11,7 @@ from mypy.types import Type as MypyType from mypy_django_plugin.django.context import DjangoContext from mypy_django_plugin.lib import fullnames, helpers -from mypy_django_plugin.transformers import fields, forms, init_create, querysets, settings, meta +from mypy_django_plugin.transformers import fields, forms, init_create, querysets, settings, meta, request from mypy_django_plugin.transformers.models import process_model_class @@ -109,7 +109,7 @@ class NewSemanalDjangoPlugin(Plugin): # for `get_user_model()` if self.django_context.settings: - if file.fullname() == 'django.contrib.auth': + if (file.fullname() == 'django.contrib.auth') or (file.fullname() in {'django.http', 'django.http.request'}): auth_user_model_name = self.django_context.settings.AUTH_USER_MODEL try: auth_user_module = self.django_context.apps_registry.get_model(auth_user_model_name).__module__ @@ -204,6 +204,11 @@ class NewSemanalDjangoPlugin(Plugin): return partial(settings.get_type_of_settings_attribute, django_context=self.django_context) + info = self._get_typeinfo_or_none(class_name) + if info and info.has_base(fullnames.HTTPREQUEST_CLASS_FULLNAME): + return partial(request.set_auth_user_model_as_type_for_request_user, django_context=self.django_context) + + # def get_type_analyze_hook(self, fullname: str # ( ): # info = self._get_typeinfo_or_none(fullname) diff --git a/mypy_django_plugin/transformers/fields.py b/mypy_django_plugin/transformers/fields.py index 016ef1b..d153700 100644 --- a/mypy_django_plugin/transformers/fields.py +++ b/mypy_django_plugin/transformers/fields.py @@ -15,9 +15,6 @@ def get_referred_to_model_fullname(ctx: FunctionContext, django_context: DjangoC return to_arg_type.ret_type.type.fullname() outer_model_info = ctx.api.scope.active_class() - # if not outer_model_info or not outer_model_info.has_base(fullnames.MODEL_CLASS_FULLNAME): - # # not inside models.Model class - # return None assert isinstance(outer_model_info, TypeInfo) to_arg_expr = helpers.get_call_argument_by_name(ctx, 'to') diff --git a/mypy_django_plugin/transformers/request.py b/mypy_django_plugin/transformers/request.py new file mode 100644 index 0000000..46bb981 --- /dev/null +++ b/mypy_django_plugin/transformers/request.py @@ -0,0 +1,15 @@ +from mypy.plugin import AttributeContext +from mypy.types import Instance, Type as MypyType + +from mypy_django_plugin.django.context import DjangoContext +from mypy_django_plugin.lib import helpers + + +def set_auth_user_model_as_type_for_request_user(ctx: AttributeContext, django_context: DjangoContext) -> MypyType: + auth_user_model = django_context.settings.AUTH_USER_MODEL + model_cls = django_context.apps_registry.get_model(auth_user_model) + model_info = helpers.lookup_class_typeinfo(ctx.api, model_cls) + if model_info is None: + return ctx.default_attr_type + + return Instance(model_info, []) diff --git a/test-data/typecheck/test_request.yml b/test-data/typecheck/test_request.yml new file mode 100644 index 0000000..e05260a --- /dev/null +++ b/test-data/typecheck/test_request.yml @@ -0,0 +1,16 @@ +- case: request_object_has_user_of_type_auth_user_model + disable_cache: true + main: | + from django.http.request import HttpRequest + reveal_type(HttpRequest().user) # N: Revealed type is 'myapp.models.MyUser' + installed_apps: + - myapp + additional_settings: + - AUTH_USER_MODEL='myapp.MyUser' + files: + - path: myapp/__init__.py + - path: myapp/models.py + content: | + from django.db import models + class MyUser(models.Model): + pass \ No newline at end of file