diff --git a/mypy_django_plugin/django/context.py b/mypy_django_plugin/django/context.py index f5ea788..7bbc8e0 100644 --- a/mypy_django_plugin/django/context.py +++ b/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): diff --git a/mypy_django_plugin/lib/fullnames.py b/mypy_django_plugin/lib/fullnames.py index 8e0a2a0..96a6219 100644 --- a/mypy_django_plugin/lib/fullnames.py +++ b/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" diff --git a/mypy_django_plugin/transformers/fields.py b/mypy_django_plugin/transformers/fields.py index 331a3f9..0827ff4 100644 --- a/mypy_django_plugin/transformers/fields.py +++ b/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: - args.append(helpers.convert_any_to_type(default_arg, base_type)) + for new_type, default_arg in zip(base_field_arg_type.args, default_return_type.args): + # 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) diff --git a/mypy_django_plugin/transformers/orm_lookups.py b/mypy_django_plugin/transformers/orm_lookups.py index 67fbd1c..acba050 100644 --- a/mypy_django_plugin/transformers/orm_lookups.py +++ b/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) diff --git a/tests/typecheck/contrib/postgres/test_fields.yml b/tests/typecheck/contrib/postgres/test_fields.yml new file mode 100644 index 0000000..5dc93cf --- /dev/null +++ b/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()) diff --git a/tests/typecheck/fields/test_base.yml b/tests/typecheck/fields/test_base.yml index c32db55..f9c671f 100644 --- a/tests/typecheck/fields/test_base.yml +++ b/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*" diff --git a/tests/typecheck/fields/test_related.yml b/tests/typecheck/fields/test_related.yml index 1c635a2..439820c 100644 --- a/tests/typecheck/fields/test_related.yml +++ b/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: diff --git a/tests/typecheck/models/test_create.yml b/tests/typecheck/models/test_create.yml index 76a50cc..3376adf 100644 --- a/tests/typecheck/models/test_create.yml +++ b/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) - MyModel.objects.create(id=None) + first = MyModel(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) diff --git a/tests/typecheck/models/test_init.yml b/tests/typecheck/models/test_init.yml index c0e62d8..2fcf84c 100644 --- a/tests/typecheck/models/test_init.yml +++ b/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=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=[], 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]") 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[]", 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())