create(id=None) is valid, if id is AutoField

mypy_django_plugin/django/context.py

@@ -8,7 +8,7 @@ from typing import (
 from django.contrib.postgres.fields import ArrayField
 from django.core.exceptions import FieldError
 from django.db.models.base import Model
-from django.db.models.fields import CharField, Field
+from django.db.models.fields import CharField, Field, AutoField
 from django.db.models.fields.related import ForeignKey, RelatedField
 from django.db.models.fields.reverse_related import ForeignObjectRel
 from django.db.models.sql.query import Query
@@ -79,6 +79,9 @@ class DjangoFieldsContext:
         if method == '__init__':
             if field.primary_key or isinstance(field, ForeignKey):
                 return True
+        if method == 'create':
+            if isinstance(field, AutoField):
+                return True
         if field.has_default():
             return True
         return nullable

mypy_django_plugin/transformers/fields.py

@@ -1,75 +1,45 @@
 from typing import Optional, Tuple, cast
 
-from mypy.nodes import MypyFile, TypeInfo
+from django.db.models.fields.related import RelatedField
+from mypy.nodes import AssignmentStmt, TypeInfo
 from mypy.plugin import FunctionContext
-from mypy.types import AnyType, CallableType, Instance
-from mypy.types import Type as MypyType
-from mypy.types import TypeOfAny
+from mypy.types import AnyType, Instance, Type as MypyType, TypeOfAny
 
+from django.db.models.fields import Field
 from mypy_django_plugin.django.context import DjangoContext
 from mypy_django_plugin.lib import fullnames, helpers
 
 
-def get_referred_to_model_fullname(ctx: FunctionContext, django_context: DjangoContext) -> Optional[str]:
-    to_arg_type = helpers.get_call_argument_type_by_name(ctx, 'to')
-    if isinstance(to_arg_type, CallableType):
-        assert isinstance(to_arg_type.ret_type, Instance)
-        return to_arg_type.ret_type.type.fullname()
-
+def _get_current_field_from_assignment(ctx: FunctionContext, django_context: DjangoContext) -> Optional[Field]:
     outer_model_info = ctx.api.scope.active_class()
     assert isinstance(outer_model_info, TypeInfo)
-
-    to_arg_expr = helpers.get_call_argument_by_name(ctx, 'to')
-    model_string = helpers.resolve_string_attribute_value(to_arg_expr, ctx, django_context)
-    if model_string is None:
-        # unresolvable
+    if not outer_model_info.has_base(fullnames.MODEL_CLASS_FULLNAME):
         return None
 
-    if model_string == 'self':
-        return outer_model_info.fullname()
-    if '.' not in model_string:
-        # same file class
-        model_cls_is_accessible = False
-        for scope in ctx.api.scope.stack:
-            if isinstance(scope, (MypyFile, TypeInfo)):
-                model_class_candidate = scope.names.get(model_string)
-                model_cls_is_accessible = (model_class_candidate is not None
-                                           and isinstance(model_class_candidate.node, TypeInfo)
-                                           and model_class_candidate.node.has_base(fullnames.MODEL_CLASS_FULLNAME))
-                if model_cls_is_accessible:
+    field_name = None
+    for stmt in outer_model_info.defn.defs.body:
+        if isinstance(stmt, AssignmentStmt):
+            if stmt.rvalue == ctx.context:
+                field_name = stmt.lvalues[0].name
                 break
-            # TODO: FuncItem
-
-        if not model_cls_is_accessible:
-            ctx.api.fail(f'No model {model_string!r} defined in the current module', ctx.context)
+    if field_name is None:
         return None
 
-        return outer_model_info.module_name + '.' + model_string
-
-    app_label, model_name = model_string.split('.')
-    if app_label not in django_context.apps_registry.app_configs:
-        ctx.api.fail(f'No installed app with label {app_label!r}', ctx.context)
+    model_cls = django_context.get_model_class_by_fullname(outer_model_info.fullname())
+    if model_cls is None:
         return None
 
-    try:
-        model_cls = django_context.apps_registry.get_model(app_label, model_name)
-    except LookupError as exc:
-        # no model in app
-        ctx.api.fail(exc.args[0], ctx.context)
-        return None
-
-    model_fullname = helpers.get_class_fullname(model_cls)
-    return model_fullname
+    current_field = model_cls._meta.get_field(field_name)
+    return current_field
 
 
 def fill_descriptor_types_for_related_field(ctx: FunctionContext, django_context: DjangoContext) -> MypyType:
-    referred_to_fullname = get_referred_to_model_fullname(ctx, django_context)
-    if referred_to_fullname is None:
+    current_field = _get_current_field_from_assignment(ctx, django_context)
+    if current_field is None:
         return AnyType(TypeOfAny.from_error)
 
-    referred_to_typeinfo = helpers.lookup_fully_qualified_generic(referred_to_fullname, ctx.api.modules)
-    assert isinstance(referred_to_typeinfo, TypeInfo), f'Cannot resolve {referred_to_fullname!r}'
-
+    assert isinstance(current_field, RelatedField)
+    referred_to_typeinfo = helpers.lookup_class_typeinfo(ctx.api, current_field.related_model)
     referred_to_type = Instance(referred_to_typeinfo, [])
 
     default_related_field_type = set_descriptor_types_for_field(ctx)

test-data/typecheck/fields/test_related.yml

@@ -478,3 +478,46 @@
                 class MyApp2Config(AppConfig):
                     name = 'myapp2'
                     label = 'myapp2__user'
+
+
+-   case: related_field_to_extracted_from_function
+    main: |
+        from myapp.models import Profile
+        reveal_type(Profile().user)  # N: Revealed type is 'myapp.models.User*'
+    installed_apps:
+        - myapp
+    files:
+        -   path: myapp/__init__.py
+        -   path: myapp/models.py
+            content: |
+                from django.db import models
+                class User(models.Model):
+                    pass
+                def get_user_model_name():
+                    return 'myapp.User'
+                class Profile(models.Model):
+                    user = models.ForeignKey(to=get_user_model_name(), on_delete=models.CASCADE)
+
+
+-   case: related_manager_name_defined_by_pattern
+    main: |
+        from myapp.models import Publisher
+        reveal_type(Publisher().books)  # N: Revealed type is 'django.db.models.manager.RelatedManager[myapp.models.Book]'
+        reveal_type(Publisher().articles)  # N: Revealed type is 'django.db.models.manager.RelatedManager[myapp.models.Article]'
+    installed_apps:
+        - myapp
+    files:
+        -   path: myapp/__init__.py
+        -   path: myapp/models.py
+            content: |
+                from django.db import models
+                class Publisher(models.Model):
+                    pass
+                class Entry(models.Model):
+                    class Meta:
+                        abstract = True
+                    publisher = models.ForeignKey(to=Publisher, related_name='%(class)ss', on_delete=models.CASCADE)
+                class Book(Entry):
+                    pass
+                class Article(Entry):
+                    pass

test-data/typecheck/models/test_create.yml

@@ -55,10 +55,11 @@
                 class Child4(Child1):
                     value4 = models.IntegerField()
 
--   case: optional_id_fields_for_create_is_error
+-   case: optional_id_fields_for_create_is_error_if_not_autofield
     main: |
         from myapp.models import Publisher, Book
-        Book.objects.create(id=None)  # E: Incompatible type for "id" of "Book" (got "None", expected "Union[Combinable, int, str]")
+
+        Book.objects.create(id=None)  # E: Incompatible type for "id" of "Book" (got "None", expected "Union[float, int, str, Combinable]")
         Book.objects.create(publisher=None)  # E: Incompatible type for "publisher" of "Book" (got "None", expected "Union[Publisher, Combinable]")
         Book.objects.create(publisher_id=None)  # E: Incompatible type for "publisher_id" of "Book" (got "None", expected "Union[Combinable, int, str]")
     installed_apps:
@@ -71,8 +72,23 @@
                 class Publisher(models.Model):
                     pass
                 class Book(models.Model):
+                    id = models.IntegerField(primary_key=True)
                     publisher = models.ForeignKey(Publisher, on_delete=models.CASCADE)
 
+-   case: none_for_primary_key_is_allowed_if_field_is_autogenerated
+    main: |
+        from myapp.models import Book
+        Book.objects.create(id=None)
+    installed_apps:
+        - myapp
+    files:
+        -   path: myapp/__init__.py
+        -   path: myapp/models.py
+            content: |
+                from django.db import models
+                class Book(models.Model):
+                    pass
+
 -   case: when_default_for_primary_key_is_specified_allow_none_to_be_set
     main: |
         from myapp.models import MyModel