From 57796077c6525fd6d987a7640a17d73104603347 Mon Sep 17 00:00:00 2001 From: Maxim Kurnikov Date: Mon, 22 Jul 2019 20:14:59 +0300 Subject: [PATCH] create(id=None) is valid, if id is AutoField --- mypy_django_plugin/django/context.py | 5 +- mypy_django_plugin/transformers/fields.py | 72 ++++++--------------- test-data/typecheck/fields/test_related.yml | 43 ++++++++++++ test-data/typecheck/models/test_create.yml | 20 +++++- 4 files changed, 86 insertions(+), 54 deletions(-) diff --git a/mypy_django_plugin/django/context.py b/mypy_django_plugin/django/context.py index fe049fd..cdcd707 100644 --- a/mypy_django_plugin/django/context.py +++ b/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 diff --git a/mypy_django_plugin/transformers/fields.py b/mypy_django_plugin/transformers/fields.py index e67dec4..f393ea9 100644 --- a/mypy_django_plugin/transformers/fields.py +++ b/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: - 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) - 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) + 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 + if field_name 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) + model_cls = django_context.get_model_class_by_fullname(outer_model_info.fullname()) + if model_cls is None: 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) diff --git a/test-data/typecheck/fields/test_related.yml b/test-data/typecheck/fields/test_related.yml index 7287166..b0bd874 100644 --- a/test-data/typecheck/fields/test_related.yml +++ b/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 \ No newline at end of file diff --git a/test-data/typecheck/models/test_create.yml b/test-data/typecheck/models/test_create.yml index 5b24716..37d46bd 100644 --- a/test-data/typecheck/models/test_create.yml +++ b/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