From 9beb5327de662569dab313757876f892746e7904 Mon Sep 17 00:00:00 2001 From: Konstantin Alekseev Date: Wed, 7 Apr 2021 11:37:28 +0300 Subject: [PATCH] Create related managers from generated managers (#580) --- mypy_django_plugin/transformers/models.py | 24 ++++++++++- scripts/enabled_test_modules.py | 2 +- tests/typecheck/fields/test_related.yml | 52 ++++++++++++++++++++++- 3 files changed, 74 insertions(+), 4 deletions(-) diff --git a/mypy_django_plugin/transformers/models.py b/mypy_django_plugin/transformers/models.py index 751f895..14d09f3 100644 --- a/mypy_django_plugin/transformers/models.py +++ b/mypy_django_plugin/transformers/models.py @@ -275,6 +275,8 @@ class AddRelatedManagers(ModelClassInitializer): # 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: + default_manager_type = self.try_generate_related_manager(related_model_cls, related_model_info) if ( default_manager_type is None or not isinstance(default_manager_type, Instance) @@ -283,12 +285,32 @@ class AddRelatedManagers(ModelClassInitializer): self.add_new_node_to_model_class(attname, parametrized_related_manager_type) continue - name = related_model_cls.__name__ + "_" + "RelatedManager" + name = model_cls.__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, [])) + def get_generated_manager_mappings(self, base_manager_fullname: str) -> Dict[str, str]: + base_manager_info = self.lookup_typeinfo(base_manager_fullname) + if base_manager_info is None or "from_queryset_managers" not in base_manager_info.metadata: + return {} + return base_manager_info.metadata["from_queryset_managers"] + + def try_generate_related_manager( + self, related_model_cls: Type[Model], related_model_info: TypeInfo + ) -> Optional[Instance]: + manager = related_model_cls._meta.managers_map["objects"] + base_manager_fullname = helpers.get_class_fullname(manager.__class__.__bases__[0]) + manager_fullname = helpers.get_class_fullname(manager.__class__) + generated_managers = self.get_generated_manager_mappings(base_manager_fullname) + if manager_fullname in generated_managers: + real_manager_fullname = generated_managers[manager_fullname] + manager_info = self.lookup_typeinfo(real_manager_fullname) # type: ignore + if manager_info: + return Instance(manager_info, [Instance(related_model_info, [])]) + return None + class AddExtraFieldMethods(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 7c049b8..3ff2cdb 100644 --- a/scripts/enabled_test_modules.py +++ b/scripts/enabled_test_modules.py @@ -290,7 +290,7 @@ IGNORED_ERRORS = { 'Incompatible types in assignment (expression has type "HttpResponseBase", variable has type "HttpResponse")', ], "many_to_many": [ - '(expression has type "List[Article]", variable has type "Article_RelatedManager2', + '(expression has type "List[Article]", variable has type "Publication_Article_RelatedManager1', '"add" of "RelatedManager" has incompatible type "Article"; expected "Union[Publication, int]"', ], "many_to_one": [ diff --git a/tests/typecheck/fields/test_related.yml b/tests/typecheck/fields/test_related.yml index 714242f..20a0450 100644 --- a/tests/typecheck/fields/test_related.yml +++ b/tests/typecheck/fields/test_related.yml @@ -652,10 +652,18 @@ - 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' + from myapp.models import User, Order, Product + reveal_type(User().orders) # N: Revealed type is 'myapp.models.User_Order_RelatedManager1' reveal_type(User().orders.get()) # N: Revealed type is 'myapp.models.Order*' reveal_type(User().orders.manager_method()) # N: Revealed type is 'builtins.int' + reveal_type(Order().products) # N: Revealed type is 'myapp.models.Order_Product_RelatedManager1' + reveal_type(Order().products.get()) # N: Revealed type is 'myapp.models.Product*' + reveal_type(Order().products.queryset_method()) # N: Revealed type is 'builtins.int' + # TODO: realted manager support to use the same type for all related managers + if 1 == 2: + manager = User().products + else: + manager = Order().products # E: Incompatible types in assignment (expression has type "Order_Product_RelatedManager1", variable has type "User_Product_RelatedManager1") installed_apps: - myapp files: @@ -671,6 +679,46 @@ class Order(models.Model): objects = OrderManager() user = models.ForeignKey(to=User, on_delete=models.CASCADE, related_name='orders') + class ProductQueryset(models.QuerySet): + def queryset_method(self) -> int: + pass + ProductManager = models.Manager.from_queryset(ProductQueryset) + class Product(models.Model): + objects = ProductManager() + order = models.ForeignKey(to=Order, on_delete=models.CASCADE, related_name='products') + user = models.ForeignKey(to=User, on_delete=models.CASCADE, related_name='products') + +- case: related_manager_no_conflict_from_star_import + main: | + import myapp.models + installed_apps: + - myapp + files: + - path: myapp/__init__.py + - path: myapp/models/__init__.py + content: | + from myapp.models.a import * + # make sure generated related manager from address to user doesn't have + # the same name with related manager from profile to user + from myapp.models.b import * + - path: myapp/models/a.py + content: | + from django.db import models + class Address(models.Model): + pass + - path: myapp/models/b.py + content: | + from django.db import models + from .a import Address + class Profile(models.Model): + pass + class UserQuerySet(models.QuerySet): + pass + UserManager = models.Manager.from_queryset(UserQuerySet) + class User(models.Model): + address = models.ForeignKey(Address, on_delete=models.CASCADE) + profile = models.ForeignKey(Profile, on_delete=models.CASCADE) + objects = UserManager() - case: many_to_many_field_can_be_used_in_alias main: |