Model.__init__ supporting same typing as assigment (#835)

* `Model.__init__` supporting same typing as assigment

* Update mypy_django_plugin/django/context.py

mypy_django_plugin/django/context.py

@@ -15,6 +15,7 @@ from django.db.models.lookups import Exact
 from django.db.models.sql.query import Query
 from django.utils.functional import cached_property
 from mypy.checker import TypeChecker
+from mypy.nodes import TypeInfo
 from mypy.plugin import MethodContext
 from mypy.types import AnyType, Instance
 from mypy.types import Type as MypyType
@@ -174,10 +175,29 @@ class DjangoContext:
             field_set_type = self.get_field_set_type(api, primary_key_field, method=method)
             expected_types["pk"] = field_set_type
 
+        def get_field_set_type_from_model_type_info(info: Optional[TypeInfo], field_name: str) -> Optional[MypyType]:
+            if info is None:
+                return None
+            field_node = info.names.get(field_name)
+            if field_node is None or not isinstance(field_node.type, Instance):
+                return None
+            elif not field_node.type.args:
+                # Field declares a set and a get type arg. Fallback to `None` when we can't find any args
+                return None
+
+            set_type = field_node.type.args[0]
+            return set_type
+
+        model_info = helpers.lookup_class_typeinfo(api, model_cls)
         for field in model_cls._meta.get_fields():
             if isinstance(field, Field):
                 field_name = field.attname
-                field_set_type = self.get_field_set_type(api, field, method=method)
+                # Try to retrieve set type from a model's TypeInfo object and fallback to retrieving it manually
+                # from django-stubs own declaration. This is to align with the setter types declared for
+                # assignment.
+                field_set_type = get_field_set_type_from_model_type_info(
+                    model_info, field_name
+                ) or self.get_field_set_type(api, field, method=method)
                 expected_types[field_name] = field_set_type
 
                 if isinstance(field, ForeignKey):

mypy_django_plugin/lib/fullnames.py

@@ -36,6 +36,7 @@ MIGRATION_CLASS_FULLNAME = "django.db.migrations.migration.Migration"
 OPTIONS_CLASS_FULLNAME = "django.db.models.options.Options"
 HTTPREQUEST_CLASS_FULLNAME = "django.http.request.HttpRequest"
 
+COMBINABLE_EXPRESSION_FULLNAME = "django.db.models.expressions.Combinable"
 F_EXPRESSION_FULLNAME = "django.db.models.expressions.F"
 
 ANY_ATTR_ALLOWED_CLASS_FULLNAME = "django_stubs_ext.AnyAttrAllowed"

mypy_django_plugin/transformers/fields.py

@@ -6,7 +6,7 @@ from mypy.nodes import AssignmentStmt, NameExpr, TypeInfo
 from mypy.plugin import FunctionContext
 from mypy.types import AnyType, Instance
 from mypy.types import Type as MypyType
-from mypy.types import TypeOfAny
+from mypy.types import TypeOfAny, UnionType
 
 from mypy_django_plugin.django.context import DjangoContext
 from mypy_django_plugin.lib import fullnames, helpers
@@ -125,6 +125,11 @@ def set_descriptor_types_for_field(
     null_expr = helpers.get_call_argument_by_name(ctx, "null")
     if null_expr is not None:
         is_nullable = helpers.parse_bool(null_expr) or False
+    # Allow setting field value to `None` when a field is primary key and has a default that can produce a value
+    default_expr = helpers.get_call_argument_by_name(ctx, "default")
+    primary_key_expr = helpers.get_call_argument_by_name(ctx, "primary_key")
+    if default_expr is not None and primary_key_expr is not None:
+        is_set_nullable = helpers.parse_bool(primary_key_expr) or False
 
     set_type, get_type = get_field_descriptor_types(
         default_return_type.type,
@@ -141,10 +146,46 @@ def determine_type_of_array_field(ctx: FunctionContext, django_context: DjangoCo
     if not base_field_arg_type or not isinstance(base_field_arg_type, Instance):
         return default_return_type
 
-    base_type = base_field_arg_type.args[1]  # extract __get__ type
+    def drop_combinable(_type: MypyType) -> Optional[MypyType]:
+        if isinstance(_type, Instance) and _type.type.has_base(fullnames.COMBINABLE_EXPRESSION_FULLNAME):
+            return None
+        elif isinstance(_type, UnionType):
+            items_without_combinable = []
+            for item in _type.items:
+                reduced = drop_combinable(item)
+                if reduced is not None:
+                    items_without_combinable.append(reduced)
+
+            if len(items_without_combinable) > 1:
+                return UnionType(
+                    items_without_combinable,
+                    line=_type.line,
+                    column=_type.column,
+                    is_evaluated=_type.is_evaluated,
+                    uses_pep604_syntax=_type.uses_pep604_syntax,
+                )
+            elif len(items_without_combinable) == 1:
+                return items_without_combinable[0]
+            else:
+                return None
+
+        return _type
+
+    # Both base_field and return type should derive from Field and thus expect 2 arguments
+    assert len(base_field_arg_type.args) == len(default_return_type.args) == 2
     args = []
-    for default_arg in default_return_type.args:
+    for new_type, default_arg in zip(base_field_arg_type.args, default_return_type.args):
-        args.append(helpers.convert_any_to_type(default_arg, base_type))
+        # Drop any base_field Combinable type
+        reduced = drop_combinable(new_type)
+        if reduced is None:
+            ctx.api.fail(
+                f"Can't have ArrayField expecting {fullnames.COMBINABLE_EXPRESSION_FULLNAME!r} as data type",
+                ctx.context,
+            )
+        else:
+            new_type = reduced
+
+        args.append(helpers.convert_any_to_type(default_arg, new_type))
 
     return helpers.reparametrize_instance(default_return_type, args)
 

mypy_django_plugin/transformers/orm_lookups.py

@@ -29,7 +29,7 @@ def typecheck_queryset_filter(ctx: MethodContext, django_context: DjangoContext)
         if lookup_kwarg is None:
             continue
         if isinstance(provided_type, Instance) and provided_type.type.has_base(
-            "django.db.models.expressions.Combinable"
+            fullnames.COMBINABLE_EXPRESSION_FULLNAME
         ):
             provided_type = resolve_combinable_type(provided_type, django_context)
 

tests/typecheck/contrib/postgres/test_fields.yml

@@ -0,0 +1,19 @@
+-   case: union_combinable_reduced_to_non_union
+    main: |
+        from typing import List
+        from myapp.models import MyModel
+        array_val: List[int] = [1]
+        MyModel(array=array_val)  # E: Incompatible type for "array" of "MyModel" (got "List[int]", expected "Union[Sequence[str], Combinable]")
+        non_init = MyModel()
+        non_init.array = array_val  # E: Incompatible types in assignment (expression has type "List[int]", variable has type "Union[Sequence[str], Combinable]")
+    installed_apps:
+        - myapp
+    files:
+        -   path: myapp/__init__.py
+        -   path: myapp/models.py
+            content: |
+                from django.db import models
+                from django.contrib.postgres.fields import ArrayField
+
+                class MyModel(models.Model):
+                    array = ArrayField(base_field=models.TextField())

tests/typecheck/fields/test_base.yml

@@ -94,7 +94,7 @@
 -   case: blank_and_not_null_charfield_does_not_allow_none
     main: |
         from myapp.models import MyModel
-        MyModel(notnulltext=None)  # Should allow None in constructor
+        MyModel(notnulltext=None)  # E: Incompatible type for "notnulltext" of "MyModel" (got "None", expected "Union[str, int, Combinable]")
         MyModel(notnulltext="")
         MyModel().notnulltext = None  # E: Incompatible types in assignment (expression has type "None", variable has type "Union[str, int, Combinable]")
         reveal_type(MyModel().notnulltext)  # N: Revealed type is "builtins.str*"

tests/typecheck/fields/test_related.yml

@@ -357,13 +357,13 @@
         reveal_type(Book().publisher_id)  # N: Revealed type is "builtins.str*"
         Book(publisher_id=1)
         Book(publisher_id='hello')
-        Book(publisher_id=datetime.datetime.now())  # E: Incompatible type for "publisher_id" of "Book" (got "datetime", expected "Union[str, int, Combinable, None]")
+        Book(publisher_id=datetime.datetime.now())  # E: Incompatible type for "publisher_id" of "Book" (got "datetime", expected "Union[str, int, Combinable]")
         Book.objects.create(publisher_id=1)
         Book.objects.create(publisher_id='hello')
 
         reveal_type(Book2().publisher_id)  # N: Revealed type is "builtins.int*"
         Book2(publisher_id=1)
-        Book2(publisher_id=[])  # E: Incompatible type for "publisher_id" of "Book2" (got "List[Any]", expected "Union[float, int, str, Combinable, None]")
+        Book2(publisher_id=[])  # E: Incompatible type for "publisher_id" of "Book2" (got "List[Any]", expected "Union[float, int, str, Combinable]")
         Book2.objects.create(publisher_id=1)
         Book2.objects.create(publisher_id=[])  # E: Incompatible type for "publisher_id" of "Book2" (got "List[Any]", expected "Union[float, int, str, Combinable]")
     installed_apps:

tests/typecheck/models/test_create.yml

@@ -95,12 +95,28 @@
 -   case: when_default_for_primary_key_is_specified_allow_none_to_be_set
     main: |
         from myapp.models import MyModel
-        MyModel(id=None)
+        first = MyModel(id=None)
-        MyModel.objects.create(id=None)
+        reveal_type(first.id)  # N: Revealed type is "builtins.int*"
+        first = MyModel.objects.create(id=None)
+        reveal_type(first.id)  # N: Revealed type is "builtins.int*"
+        first = MyModel()
+        first.id = None
+        reveal_type(first.id)  # N: Revealed type is "builtins.int*"
 
         from myapp.models import MyModel2
-        MyModel2(id=None)
+        MyModel2(id=None)  # E: Incompatible type for "id" of "MyModel2" (got "None", expected "Union[float, int, str, Combinable]")
         MyModel2.objects.create(id=None)  # E: Incompatible type for "id" of "MyModel2" (got "None", expected "Union[float, int, str, Combinable]")
+        second = MyModel2()
+        second.id = None  # E: Incompatible types in assignment (expression has type "None", variable has type "Union[float, int, str, Combinable]")
+        reveal_type(second.id)  # N: Revealed type is "builtins.int*"
+
+        # default set but no primary key doesn't allow None
+        from myapp.models import MyModel3
+        MyModel3(default=None)  # E: Incompatible type for "default" of "MyModel3" (got "None", expected "Union[float, int, str, Combinable]")
+        MyModel3.objects.create(default=None)  # E: Incompatible type for "default" of "MyModel3" (got "None", expected "Union[float, int, str, Combinable]")
+        third = MyModel3()
+        third.default = None  # E: Incompatible types in assignment (expression has type "None", variable has type "Union[float, int, str, Combinable]")
+        reveal_type(third.default)  # N: Revealed type is "builtins.int*"
     installed_apps:
         - myapp
     files:
@@ -114,3 +130,5 @@
                     id = models.IntegerField(primary_key=True, default=return_int)
                 class MyModel2(models.Model):
                     id = models.IntegerField(primary_key=True)
+                class MyModel3(models.Model):
+                    default = models.IntegerField(default=return_int)

tests/typecheck/models/test_init.yml

@@ -126,8 +126,8 @@
 
         from myapp.models import Publisher, PublisherDatetime, Book
         Book(publisher_id=1, publisher_dt_id=now)
-        Book(publisher_id=[], publisher_dt_id=now)  # E: Incompatible type for "publisher_id" of "Book" (got "List[Any]", expected "Union[Combinable, int, str, None]")
+        Book(publisher_id=[], publisher_dt_id=now)  # E: Incompatible type for "publisher_id" of "Book" (got "List[Any]", expected "Union[Combinable, int, str]")
-        Book(publisher_id=1, publisher_dt_id=1)  # E: Incompatible type for "publisher_dt_id" of "Book" (got "int", expected "Union[str, datetime, date, Combinable, None]")
+        Book(publisher_id=1, publisher_dt_id=1)  # E: Incompatible type for "publisher_dt_id" of "Book" (got "int", expected "Union[str, datetime, date, Combinable]")
     installed_apps:
         - myapp
     files:
@@ -155,15 +155,17 @@
         class NotAValid:
             pass
         array_val3: List[NotAValid] = [NotAValid()]
-        MyModel(array=array_val3)  # E: Incompatible type for "array" of "MyModel" (got "List[NotAValid]", expected "Union[Sequence[Union[float, int, str, Combinable]], Combinable]")
+        MyModel(array=array_val3)  # E: Incompatible type for "array" of "MyModel" (got "List[NotAValid]", expected "Union[Sequence[Union[float, int, str]], Combinable]")
+        non_init = MyModel()
+        non_init.array = array_val
+        non_init.array = array_val2
+        non_init.array = array_val3  # E: Incompatible types in assignment (expression has type "List[NotAValid]", variable has type "Union[Sequence[Union[float, int, str]], Combinable]")
     installed_apps:
         - myapp
     files:
         -   path: myapp/__init__.py
         -   path: myapp/models.py
             content: |
-                from typing import List, Tuple
-
                 from django.db import models
                 from django.contrib.postgres.fields import ArrayField
 
@@ -232,7 +234,7 @@
                 class Publisher(models.Model):
                     name = models.CharField(primary_key=True, max_length=100)
                 class Book(models.Model):
-                    publisher = models.ForeignKey(Publisher, on_delete=models.CASCADE)
+                    publisher = models.ForeignKey(Publisher, on_delete=models.CASCADE, null=True)
 
 
 -   case: init_in_abstract_model_classmethod_should_not_throw_error_for_valid_fields
@@ -255,3 +257,61 @@
                         return cls(text='mytext')
                 class MyModel(AbstractModel):
                     pass
+
+
+-   case: field_set_type_honors_type_redefinition
+    main: |
+        from typing import List
+        from myapp.models import MyModel
+        non_init = MyModel()
+        reveal_type(non_init.redefined_set_type)
+        reveal_type(non_init.redefined_union_set_type)
+        reveal_type(non_init.redefined_array_set_type)
+        reveal_type(non_init.default_set_type)
+        reveal_type(non_init.unset_set_type)
+        non_init.redefined_set_type = "invalid"
+        non_init.redefined_union_set_type = "invalid"
+        array_val: List[str] = ["invalid"]
+        non_init.redefined_array_set_type = array_val
+        non_init.default_set_type = []
+        non_init.unset_set_type = []
+        MyModel(
+            redefined_set_type="invalid",
+            redefined_union_set_type="invalid",
+            redefined_array_set_type=33,
+            default_set_type=[],
+            unset_set_type=[],
+        )
+    out: |
+        main:4: note: Revealed type is "builtins.int*"
+        main:5: note: Revealed type is "builtins.int*"
+        main:6: note: Revealed type is "builtins.list*[builtins.int]"
+        main:7: note: Revealed type is "builtins.int*"
+        main:8: note: Revealed type is "Any"
+        main:9: error: Incompatible types in assignment (expression has type "str", variable has type "int")
+        main:10: error: Incompatible types in assignment (expression has type "str", variable has type "Union[int, float]")
+        main:12: error: Incompatible types in assignment (expression has type "List[str]", variable has type "Sequence[Union[int, float]]")
+        main:13: error: Incompatible types in assignment (expression has type "List[<nothing>]", variable has type "Union[float, int, str, Combinable]")
+        main:15: error: Incompatible type for "redefined_set_type" of "MyModel" (got "str", expected "int")
+        main:15: error: Incompatible type for "redefined_union_set_type" of "MyModel" (got "str", expected "Union[int, float]")
+        main:15: error: Incompatible type for "redefined_array_set_type" of "MyModel" (got "int", expected "Sequence[Union[int, float]]")
+        main:15: error: Incompatible type for "default_set_type" of "MyModel" (got "List[Any]", expected "Union[float, int, str, Combinable]")
+    installed_apps:
+        - myapp
+    files:
+        -   path: myapp/__init__.py
+        -   path: myapp/models.py
+            content: |
+                from django.contrib.postgres.fields import ArrayField
+                from django.db import models
+                from typing import cast, List, Sequence, Union
+
+                class MyModel(models.Model):
+                    redefined_set_type = cast("models.Field[int, int]", models.IntegerField())
+                    redefined_union_set_type = cast("models.Field[Union[int, float], int]", models.IntegerField())
+                    redefined_array_set_type = cast(
+                        "ArrayField[Sequence[Union[int, float]], List[int]]",
+                        ArrayField(base_field=models.IntegerField()),
+                    )
+                    default_set_type = models.IntegerField()
+                    unset_set_type = cast("models.Field", models.IntegerField())