From 2c4827bbaf6406861168ef7c7c1a0ec206a60ed0 Mon Sep 17 00:00:00 2001 From: Maxim Kurnikov Date: Sat, 30 Nov 2019 21:52:02 +0300 Subject: [PATCH] properly change type of self for methods on custom manager classes --- mypy_django_plugin/transformers/models.py | 37 +++++++++++++++++-- scripts/enabled_test_modules.py | 6 +-- .../typecheck/managers/test_managers.yml | 17 ++++++++- 3 files changed, 51 insertions(+), 9 deletions(-) diff --git a/mypy_django_plugin/transformers/models.py b/mypy_django_plugin/transformers/models.py index 7cd26b7..f85c661 100644 --- a/mypy_django_plugin/transformers/models.py +++ b/mypy_django_plugin/transformers/models.py @@ -1,5 +1,5 @@ from collections import OrderedDict -from typing import Type +from typing import List, Tuple, Type from django.db.models.base import Model from django.db.models.fields import DateField, DateTimeField @@ -7,10 +7,11 @@ from django.db.models.fields.related import ForeignKey from django.db.models.fields.reverse_related import ( ManyToManyRel, ManyToOneRel, OneToOneRel, ) -from mypy.nodes import ARG_STAR2, Argument, Context, TypeInfo, Var +from mypy.nodes import ARG_STAR2, Argument, Context, FuncDef, TypeInfo, Var from mypy.plugin import ClassDefContext from mypy.plugins import common -from mypy.types import AnyType, Instance +from mypy.plugins.common import add_method +from mypy.types import AnyType, CallableType, Instance from mypy.types import Type as MypyType from mypy.types import TypeOfAny @@ -158,7 +159,22 @@ class AddManagers(ModelClassInitializer): bases=bases, fields=OrderedDict()) # copy fields to a new manager + new_cls_def_context = ClassDefContext(cls=custom_manager_info.defn, + reason=self.ctx.reason, + api=self.api) + custom_manager_type = Instance(custom_manager_info, [Instance(self.model_classdef.info, [])]) + for name, sym in manager_info.names.items(): + # replace self type with new class, if copying method + if isinstance(sym.node, FuncDef): + arguments, return_type = self.prepare_new_method_arguments(sym.node) + add_method(new_cls_def_context, + name, + args=arguments, + return_type=return_type, + self_type=custom_manager_type) + continue + new_sym = sym.copy() if isinstance(new_sym.node, Var): new_var = Var(name, type=sym.type) @@ -167,9 +183,22 @@ class AddManagers(ModelClassInitializer): new_sym.node = new_var custom_manager_info.names[name] = new_sym - custom_manager_type = Instance(custom_manager_info, [Instance(self.model_classdef.info, [])]) self.add_new_node_to_model_class(manager_name, custom_manager_type) + def prepare_new_method_arguments(self, node: FuncDef) -> Tuple[List[Argument], MypyType]: + arguments = [] + for argument in node.arguments[1:]: + if argument.type_annotation is None: + argument.type_annotation = AnyType(TypeOfAny.unannotated) + arguments.append(argument) + + if isinstance(node.type, CallableType): + return_type = node.type.ret_type + else: + return_type = AnyType(TypeOfAny.unannotated) + + return arguments, return_type + class AddDefaultManagerAttribute(ModelClassInitializer): def run_with_model_cls(self, model_cls: Type[Model]) -> None: diff --git a/scripts/enabled_test_modules.py b/scripts/enabled_test_modules.py index 2bb30b9..7eba8b9 100644 --- a/scripts/enabled_test_modules.py +++ b/scripts/enabled_test_modules.py @@ -37,7 +37,6 @@ IGNORED_ERRORS = { 'Argument after ** must be a mapping', 'note:', re.compile(r'Item "None" of "[a-zA-Z_ ,\[\]]+" has no attribute'), - '"Optional[List[_Record]]"', '"Callable[..., None]" has no attribute', 'does not return a value', 'has no attribute "alternatives"', @@ -257,7 +256,6 @@ IGNORED_ERRORS = { 'Item "Field[Any, Any]" of "Union[Field[Any, Any], ForeignObjectRel]" has no attribute', 'Incompatible types in assignment (expression has type "Type[Person', 'Incompatible types in assignment (expression has type "FloatModel", variable has type', - 'No overload variant of "bytes" matches argument type "Container[int]"', '"ImageFile" has no attribute "was_opened"', ], 'model_indexes': [ @@ -278,6 +276,7 @@ IGNORED_ERRORS = { 'Incompatible type for "department" of "Worker"', '"PickledModel" has no attribute', '"Department" has no attribute "evaluate"', + 'Unsupported target for indexed assignment', ], 'model_formsets_regress': [ 'Incompatible types in assignment (expression has type "int", target has type "str")', @@ -352,7 +351,8 @@ IGNORED_ERRORS = { 'Argument 1 to "append" of "list" has incompatible type "Tuple[Any, Any, Optional[Any], Any]";' ], 'sites_framework': [ - 'expression has type "CurrentSiteManager[CustomArticle]", base class "AbstractArticle"' + 'expression has type "CurrentSiteManager[CustomArticle]", base class "AbstractArticle"', + "Name 'Optional' is not defined", ], 'syndication_tests': [ 'List or tuple expected as variable arguments' diff --git a/test-data/typecheck/managers/test_managers.yml b/test-data/typecheck/managers/test_managers.yml index 885ee0e..9cc8698 100644 --- a/test-data/typecheck/managers/test_managers.yml +++ b/test-data/typecheck/managers/test_managers.yml @@ -307,9 +307,18 @@ - case: custom_manager_returns_proper_model_types main: | from myapp.models import User - reveal_type(User.objects.get()) # N: Revealed type is 'myapp.models.User*' + reveal_type(User.objects) # N: Revealed type is 'myapp.models.User_MyManager[myapp.models.User]' reveal_type(User.objects.select_related()) # N: Revealed type is 'myapp.models.User_MyManager[myapp.models.User]' + reveal_type(User.objects.get()) # N: Revealed type is 'myapp.models.User*' reveal_type(User.objects.get_instance()) # N: Revealed type is 'builtins.int' + reveal_type(User.objects.get_instance_untyped('hello')) # N: Revealed type is 'Any' + + from myapp.models import ChildUser + reveal_type(ChildUser.objects) # N: Revealed type is 'myapp.models.ChildUser_MyManager[myapp.models.ChildUser]' + reveal_type(ChildUser.objects.select_related()) # N: Revealed type is 'myapp.models.ChildUser_MyManager[myapp.models.ChildUser]' + reveal_type(ChildUser.objects.get()) # N: Revealed type is 'myapp.models.ChildUser*' + reveal_type(ChildUser.objects.get_instance()) # N: Revealed type is 'builtins.int' + reveal_type(ChildUser.objects.get_instance_untyped('hello')) # N: Revealed type is 'Any' installed_apps: - myapp files: @@ -318,7 +327,11 @@ content: | from django.db import models class MyManager(models.Manager): - def get_instance(self: "models.Manager[User]") -> int: + def get_instance(self) -> int: + pass + def get_instance_untyped(self, name): pass class User(models.Model): objects = MyManager() + class ChildUser(models.Model): + objects = MyManager()