Fix suggested changes from #820 (#821)

mypy_django_plugin/main.py

@@ -309,37 +309,51 @@ class NewSemanalDjangoPlugin(Plugin):
         return None
 
     def get_base_class_hook(self, fullname: str) -> Optional[Callable[[ClassDefContext], None]]:
+        # Base class is a Model class definition
         if (
             fullname in self.django_context.all_registered_model_class_fullnames
             or fullname in self._get_current_model_bases()
         ):
             return partial(transform_model_class, django_context=self.django_context)
 
+        # Base class is a Manager class definition
         if fullname in self._get_current_manager_bases():
             return add_new_manager_base_hook
 
+        # Base class is a Form class definition
         if fullname in self._get_current_form_bases():
             return transform_form_class
         return None
 
     def get_attribute_hook(self, fullname: str) -> Optional[Callable[[AttributeContext], MypyType]]:
         class_name, _, attr_name = fullname.rpartition(".")
+
+        # Lookup of a settings variable
         if class_name == fullnames.DUMMY_SETTINGS_BASE_CLASS:
             return partial(settings.get_type_of_settings_attribute, django_context=self.django_context)
 
         info = self._get_typeinfo_or_none(class_name)
+
+        # Lookup of the '.is_superuser' attribute
         if info and info.has_base(fullnames.PERMISSION_MIXIN_CLASS_FULLNAME) and attr_name == "is_superuser":
             return partial(set_auth_user_model_boolean_fields, django_context=self.django_context)
+
+        # Lookup of the 'request.user' attribute
         if info and info.has_base(fullnames.HTTPREQUEST_CLASS_FULLNAME) and attr_name == "user":
             return partial(request.set_auth_user_model_as_type_for_request_user, django_context=self.django_context)
+
+        # Lookup of the 'user.is_staff' or 'user.is_active' attribute
         if info and info.has_base(fullnames.ABSTRACT_USER_MODEL_FULLNAME) and attr_name in ("is_staff", "is_active"):
             return partial(set_auth_user_model_boolean_fields, django_context=self.django_context)
+
+        # Lookup of a method on a dynamically generated manager class
+        # i.e. a manager class only existing while mypy is running, not collected from the AST
         if (
             info
             and info.has_base(fullnames.BASE_MANAGER_CLASS_FULLNAME)
             and class_name in self._get_current_manager_bases()
         ):
-            return partial(resolve_manager_method, django_context=self.django_context)
+            return resolve_manager_method
 
         return None
 
@@ -352,6 +366,7 @@ class NewSemanalDjangoPlugin(Plugin):
             return partial(handle_annotated_type, django_context=self.django_context)
 
     def get_dynamic_class_hook(self, fullname: str) -> Optional[Callable[[DynamicClassDefContext], None]]:
+        # Create a new manager class definition when a manager's '.from_queryset' classmethod is called
         if fullname.endswith("from_queryset"):
             class_name, _, _ = fullname.rpartition(".")
             info = self._get_typeinfo_or_none(class_name)

mypy_django_plugin/transformers/managers.py

@@ -20,7 +20,6 @@ from mypy.types import AnyType, CallableType, Instance, ProperType
 from mypy.types import Type as MypyType
 from mypy.types import TypeOfAny, TypeVarType, UnboundType, get_proper_type
 
-from mypy_django_plugin.django.context import DjangoContext
 from mypy_django_plugin.lib import fullnames, helpers
 
 
@@ -96,7 +95,7 @@ def get_method_type_from_reverse_manager(
     )
 
 
-def resolve_manager_method(ctx: AttributeContext, django_context: DjangoContext) -> MypyType:
+def resolve_manager_method(ctx: AttributeContext) -> MypyType:
     """
     A 'get_attribute_hook' that is intended to be invoked whenever the TypeChecker encounters
     an attribute on a class that has 'django.db.models.BaseManager' as a base.
@@ -179,9 +178,7 @@ def create_new_manager_class_from_from_queryset_method(ctx: DynamicClassDefConte
     new_manager_info.metaclass_type = new_manager_info.calculate_metaclass_type()
     # Stash the queryset fullname which was passed to .from_queryset
     # So that our 'resolve_manager_method' attribute hook can fetch the method from that QuerySet class
-    new_manager_info.metadata.setdefault("django", {})
-    new_manager_info.metadata["django"].setdefault("from_queryset_manager", {})
-    new_manager_info.metadata["django"]["from_queryset_manager"] = derived_queryset_fullname
+    new_manager_info.metadata["django"] = {"from_queryset_manager": derived_queryset_fullname}
 
     if len(ctx.call.args) > 1:
         expr = ctx.call.args[1]

mypy_django_plugin/transformers/models.py

@@ -283,6 +283,7 @@ class AddDefaultManagerAttribute(ModelClassInitializer):
             if not self.api.final_iteration:
                 raise exc
             else:
+                # On final round, see if the default manager is a generated (dynamic class) manager
                 base_manager_fullname = helpers.get_class_fullname(default_manager_cls.__bases__[0])
                 generated_manager_info = self.get_generated_manager_info(
                     default_manager_fullname, base_manager_fullname