From 665f4d8ea1723d0bcae59904a590ffb03a7c0355 Mon Sep 17 00:00:00 2001 From: Maksim Kurnikov Date: Tue, 17 Dec 2019 19:06:27 +0300 Subject: [PATCH] Make related manager inherit from `objects` of related model (#278) * related manager inherits from objects of related model * fix test typechecking * lint --- mypy_django_plugin/transformers/models.py | 45 +++++++++++++++------ scripts/enabled_test_modules.py | 2 +- test-data/typecheck/fields/test_related.yml | 24 +++++++++++ 3 files changed, 57 insertions(+), 14 deletions(-) diff --git a/mypy_django_plugin/transformers/models.py b/mypy_django_plugin/transformers/models.py index a095b1a..f0c436c 100644 --- a/mypy_django_plugin/transformers/models.py +++ b/mypy_django_plugin/transformers/models.py @@ -1,4 +1,4 @@ -from typing import Dict, Optional, Type, cast +from typing import Dict, List, Optional, Type, cast from django.db.models.base import Model from django.db.models.fields import DateField, DateTimeField @@ -58,6 +58,12 @@ class ModelClassInitializer: name=name, sym_type=typ) + def add_new_class_for_current_module(self, name: str, bases: List[Instance]) -> TypeInfo: + current_module = self.api.modules[self.model_classdef.info.module_name] + new_class_info = helpers.add_new_class_for_module(current_module, + name=name, bases=bases) + return new_class_info + def run(self) -> None: model_cls = self.django_context.get_model_class_by_fullname(self.model_classdef.fullname) if model_cls is None: @@ -164,14 +170,12 @@ class AddManagers(ModelClassInitializer): [Instance(self.model_classdef.info, [])]) bases.append(original_base) - current_module = self.api.modules[self.model_classdef.info.module_name] - custom_manager_info = helpers.add_new_class_for_module(current_module, - name=name, bases=bases) + new_manager_info = self.add_new_class_for_current_module(name, bases) # copy fields to a new manager - new_cls_def_context = ClassDefContext(cls=custom_manager_info.defn, + new_cls_def_context = ClassDefContext(cls=new_manager_info.defn, reason=self.ctx.reason, api=self.api) - custom_manager_type = Instance(custom_manager_info, [Instance(self.model_classdef.info, [])]) + custom_manager_type = Instance(new_manager_info, [Instance(self.model_classdef.info, [])]) for name, sym in base_manager_info.names.items(): # replace self type with new class, if copying method @@ -185,10 +189,10 @@ class AddManagers(ModelClassInitializer): new_sym = sym.copy() if isinstance(new_sym.node, Var): new_var = Var(name, type=sym.type) - new_var.info = custom_manager_info - new_var._fullname = custom_manager_info.fullname + '.' + name + new_var.info = new_manager_info + new_var._fullname = new_manager_info.fullname + '.' + name new_sym.node = new_var - custom_manager_info.names[name] = new_sym + new_manager_info.names[name] = new_sym return custom_manager_type @@ -268,15 +272,30 @@ class AddRelatedManagers(ModelClassInitializer): if isinstance(relation, (ManyToOneRel, ManyToManyRel)): try: - manager_info = self.lookup_typeinfo_or_incomplete_defn_error(fullnames.RELATED_MANAGER_CLASS) + related_manager_info = self.lookup_typeinfo_or_incomplete_defn_error(fullnames.RELATED_MANAGER_CLASS) # noqa: E501 + if 'objects' not in related_model_info.names: + raise helpers.IncompleteDefnException() except helpers.IncompleteDefnException as exc: if not self.api.final_iteration: raise exc else: continue - self.add_new_node_to_model_class(attname, - Instance(manager_info, [Instance(related_model_info, [])])) - continue + + # create new RelatedManager subclass + parametrized_related_manager_type = Instance(related_manager_info, + [Instance(related_model_info, [])]) + default_manager_type = related_model_info.names['objects'].type + if (default_manager_type is None + or not isinstance(default_manager_type, Instance) + or default_manager_type.type.fullname == fullnames.MANAGER_CLASS_FULLNAME): + self.add_new_node_to_model_class(attname, parametrized_related_manager_type) + continue + + name = related_model_cls.__name__ + '_' + 'RelatedManager' + bases = [parametrized_related_manager_type, default_manager_type] + new_related_manager_info = self.add_new_class_for_current_module(name, bases) + + self.add_new_node_to_model_class(attname, Instance(new_related_manager_info, [])) class AddExtraFieldMethods(ModelClassInitializer): diff --git a/scripts/enabled_test_modules.py b/scripts/enabled_test_modules.py index 5127bae..de05624 100644 --- a/scripts/enabled_test_modules.py +++ b/scripts/enabled_test_modules.py @@ -248,7 +248,7 @@ IGNORED_ERRORS = { re.compile(r'"(HttpRequest|WSGIRequest)" has no attribute'), ], 'many_to_many': [ - '(expression has type "List[Article]", variable has type "RelatedManager[Article]"', + '(expression has type "List[Article]", variable has type "Article_RelatedManager2', '"add" of "RelatedManager" has incompatible type "Article"; expected "Union[Publication, int]"', ], 'many_to_one': [ diff --git a/test-data/typecheck/fields/test_related.yml b/test-data/typecheck/fields/test_related.yml index 473c6b9..0a79d19 100644 --- a/test-data/typecheck/fields/test_related.yml +++ b/test-data/typecheck/fields/test_related.yml @@ -648,3 +648,27 @@ abstract = True class User(AbstractUser): pass + + +- case: related_manager_is_a_subclass_of_default_manager + main: | + from myapp.models import User + reveal_type(User().orders) # N: Revealed type is 'myapp.models.Order_RelatedManager' + reveal_type(User().orders.get()) # N: Revealed type is 'myapp.models.Order*' + reveal_type(User().orders.manager_method()) # N: Revealed type is 'builtins.int' + installed_apps: + - myapp + files: + - path: myapp/__init__.py + - path: myapp/models.py + content: | + from django.db import models + class User(models.Model): + pass + class OrderManager(models.Manager): + def manager_method(self) -> int: + pass + class Order(models.Model): + objects = OrderManager() + user = models.ForeignKey(to=User, on_delete=models.CASCADE, related_name='orders') +