Skip to content

Commit

Permalink
Resolve any settings.AUTH_USER_MODEL used as to in relation
Browse files Browse the repository at this point in the history
e.g.

```python
OneToOneField(settings.AUTH_USER_MODEL, ...)
ForeignKey(settings.AUTH_USER_MODEL, ...)
ManyToManyField(settings.AUTH_USER_MODEL, ...)
```
  • Loading branch information
flaeppe committed Sep 28, 2023
1 parent 939f91c commit 9bf6dc5
Show file tree
Hide file tree
Showing 3 changed files with 57 additions and 5 deletions.
1 change: 1 addition & 0 deletions mypy_django_plugin/lib/fullnames.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
ONETOONE_FIELD_FULLNAME = "django.db.models.fields.related.OneToOneField"
MANYTOMANY_FIELD_FULLNAME = "django.db.models.fields.related.ManyToManyField"
DUMMY_SETTINGS_BASE_CLASS = "django.conf._DjangoConfLazyObject"
AUTH_USER_MODEL_FULLNAME = "django.conf.settings.AUTH_USER_MODEL"

QUERYSET_CLASS_FULLNAME = "django.db.models.query._QuerySet"
BASE_MANAGER_CLASS_FULLNAME = "django.db.models.manager.BaseManager"
Expand Down
19 changes: 14 additions & 5 deletions mypy_django_plugin/transformers/manytomany.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
from typing import NamedTuple, Optional, Union

from mypy.checker import TypeChecker
from mypy.nodes import AssignmentStmt, Expression, NameExpr, StrExpr, TypeInfo
from mypy.nodes import AssignmentStmt, Expression, MemberExpr, NameExpr, StrExpr, TypeInfo
from mypy.plugin import FunctionContext
from mypy.semanal import SemanticAnalyzer
from mypy.types import Instance, ProperType, UninhabitedType
Expand Down Expand Up @@ -129,16 +129,25 @@ def get_model_from_expression(
Attempts to resolve an expression to a 'TypeInfo' instance. Any lazy reference
argument(e.g. "<app_label>.<object_name>") to a Django model is also attempted.
"""
# TODO: Handle settings.AUTH_USER_MODEL?
if isinstance(expr, NameExpr) and isinstance(expr.node, TypeInfo):
if (
expr.node.metaclass_type is not None
and expr.node.metaclass_type.type.fullname == fullnames.MODEL_METACLASS_FULLNAME
):
return Instance(expr.node, [])
elif isinstance(expr, StrExpr):
model_info = helpers.resolve_lazy_reference(expr.value, api=api, django_context=django_context, ctx=expr)

lazy_reference = None
if isinstance(expr, StrExpr):
lazy_reference = expr.value
elif (
isinstance(expr, MemberExpr)
and isinstance(expr.expr, NameExpr)
and f"{expr.expr.fullname}.{expr.name}" == fullnames.AUTH_USER_MODEL_FULLNAME
):
lazy_reference = django_context.settings.AUTH_USER_MODEL

if lazy_reference is not None:
model_info = helpers.resolve_lazy_reference(lazy_reference, api=api, django_context=django_context, ctx=expr)
if model_info is not None:
return Instance(model_info, [])

return None
42 changes: 42 additions & 0 deletions tests/typecheck/models/test_contrib_models.yml
Original file line number Diff line number Diff line change
Expand Up @@ -62,3 +62,45 @@
...
class MyUser(AbstractUser):
objects: ClassVar[MyUserManager] = MyUserManager()
- case: test_relation_specified_by_auth_user_model
main: |
from other.models import Other
reveal_type(Other().users.get())
reveal_type(Other.users.through)
reveal_type(Other.users.through().myuser)
reveal_type(Other.users.through.objects.get().myuser)
reveal_type(Other().user)
reveal_type(Other().unq_user)
out: |
main:2: note: Revealed type is "myapp.models.MyUser"
main:3: note: Revealed type is "Type[other.models.Other_users]"
main:4: note: Revealed type is "myapp.models.MyUser"
main:5: note: Revealed type is "myapp.models.MyUser"
main:7: note: Revealed type is "myapp.models.MyUser"
main:9: note: Revealed type is "myapp.models.MyUser"
custom_settings: |
INSTALLED_APPS = ('django.contrib.contenttypes', 'django.contrib.auth', 'myapp', 'other')
AUTH_USER_MODEL='myapp.MyUser'
files:
- path: myapp/__init__.py
- path: myapp/models.py
content: |
from django.db import models
class MyUser(models.Model):
...
- path: other/__init__.py
- path: other/models.py
content: |
from django.conf import settings
from django.db import models
class Other(models.Model):
users = models.ManyToManyField(settings.AUTH_USER_MODEL)
user = models.ForeignKey(settings.AUTH_USER_MODEL, on_delete=models.CASCADE)
unq_user = models.OneToOneField(settings.AUTH_USER_MODEL, on_delete=models.CASCADE)

0 comments on commit 9bf6dc5

Please sign in to comment.