Skip to content

Commit 9bf6dc5

Browse files
committed
Resolve any settings.AUTH_USER_MODEL used as to in relation
e.g. ```python OneToOneField(settings.AUTH_USER_MODEL, ...) ForeignKey(settings.AUTH_USER_MODEL, ...) ManyToManyField(settings.AUTH_USER_MODEL, ...) ```
1 parent 939f91c commit 9bf6dc5

File tree

3 files changed

+57
-5
lines changed

3 files changed

+57
-5
lines changed

mypy_django_plugin/lib/fullnames.py

+1
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
ONETOONE_FIELD_FULLNAME = "django.db.models.fields.related.OneToOneField"
1313
MANYTOMANY_FIELD_FULLNAME = "django.db.models.fields.related.ManyToManyField"
1414
DUMMY_SETTINGS_BASE_CLASS = "django.conf._DjangoConfLazyObject"
15+
AUTH_USER_MODEL_FULLNAME = "django.conf.settings.AUTH_USER_MODEL"
1516

1617
QUERYSET_CLASS_FULLNAME = "django.db.models.query._QuerySet"
1718
BASE_MANAGER_CLASS_FULLNAME = "django.db.models.manager.BaseManager"

mypy_django_plugin/transformers/manytomany.py

+14-5
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
from typing import NamedTuple, Optional, Union
22

33
from mypy.checker import TypeChecker
4-
from mypy.nodes import AssignmentStmt, Expression, NameExpr, StrExpr, TypeInfo
4+
from mypy.nodes import AssignmentStmt, Expression, MemberExpr, NameExpr, StrExpr, TypeInfo
55
from mypy.plugin import FunctionContext
66
from mypy.semanal import SemanticAnalyzer
77
from mypy.types import Instance, ProperType, UninhabitedType
@@ -129,16 +129,25 @@ def get_model_from_expression(
129129
Attempts to resolve an expression to a 'TypeInfo' instance. Any lazy reference
130130
argument(e.g. "<app_label>.<object_name>") to a Django model is also attempted.
131131
"""
132-
# TODO: Handle settings.AUTH_USER_MODEL?
133132
if isinstance(expr, NameExpr) and isinstance(expr.node, TypeInfo):
134133
if (
135134
expr.node.metaclass_type is not None
136135
and expr.node.metaclass_type.type.fullname == fullnames.MODEL_METACLASS_FULLNAME
137136
):
138137
return Instance(expr.node, [])
139-
elif isinstance(expr, StrExpr):
140-
model_info = helpers.resolve_lazy_reference(expr.value, api=api, django_context=django_context, ctx=expr)
138+
139+
lazy_reference = None
140+
if isinstance(expr, StrExpr):
141+
lazy_reference = expr.value
142+
elif (
143+
isinstance(expr, MemberExpr)
144+
and isinstance(expr.expr, NameExpr)
145+
and f"{expr.expr.fullname}.{expr.name}" == fullnames.AUTH_USER_MODEL_FULLNAME
146+
):
147+
lazy_reference = django_context.settings.AUTH_USER_MODEL
148+
149+
if lazy_reference is not None:
150+
model_info = helpers.resolve_lazy_reference(lazy_reference, api=api, django_context=django_context, ctx=expr)
141151
if model_info is not None:
142152
return Instance(model_info, [])
143-
144153
return None

tests/typecheck/models/test_contrib_models.yml

+42
Original file line numberDiff line numberDiff line change
@@ -62,3 +62,45 @@
6262
...
6363
class MyUser(AbstractUser):
6464
objects: ClassVar[MyUserManager] = MyUserManager()
65+
66+
- case: test_relation_specified_by_auth_user_model
67+
main: |
68+
from other.models import Other
69+
reveal_type(Other().users.get())
70+
reveal_type(Other.users.through)
71+
reveal_type(Other.users.through().myuser)
72+
reveal_type(Other.users.through.objects.get().myuser)
73+
74+
reveal_type(Other().user)
75+
76+
reveal_type(Other().unq_user)
77+
out: |
78+
main:2: note: Revealed type is "myapp.models.MyUser"
79+
main:3: note: Revealed type is "Type[other.models.Other_users]"
80+
main:4: note: Revealed type is "myapp.models.MyUser"
81+
main:5: note: Revealed type is "myapp.models.MyUser"
82+
83+
main:7: note: Revealed type is "myapp.models.MyUser"
84+
85+
main:9: note: Revealed type is "myapp.models.MyUser"
86+
custom_settings: |
87+
INSTALLED_APPS = ('django.contrib.contenttypes', 'django.contrib.auth', 'myapp', 'other')
88+
AUTH_USER_MODEL='myapp.MyUser'
89+
files:
90+
- path: myapp/__init__.py
91+
- path: myapp/models.py
92+
content: |
93+
from django.db import models
94+
95+
class MyUser(models.Model):
96+
...
97+
- path: other/__init__.py
98+
- path: other/models.py
99+
content: |
100+
from django.conf import settings
101+
from django.db import models
102+
103+
class Other(models.Model):
104+
users = models.ManyToManyField(settings.AUTH_USER_MODEL)
105+
user = models.ForeignKey(settings.AUTH_USER_MODEL, on_delete=models.CASCADE)
106+
unq_user = models.OneToOneField(settings.AUTH_USER_MODEL, on_delete=models.CASCADE)

0 commit comments

Comments
 (0)