Implement support for <QuerySet>.as_manager()

Author: flaeppe

README.md

@@ -140,9 +140,7 @@ And then use `AuthenticatedHttpRequest` instead of the standard `HttpRequest` fo
 
 ### My QuerySet methods are returning Any rather than my Model
 
-`QuerySet.as_manager()` is not currently supported.
-
-If you are using `MyQuerySet.as_manager()`, then your `Manager`/`QuerySet` methods will all not be linked to your model.
+If you are using `MyQuerySet.as_manager()`:
 
 Example:
 
@@ -156,12 +154,12 @@ class MyModel(models.Model):
   bar = models.IntegerField()
   objects = MyModelQuerySet.as_manager()
 
-def use_my_model():
-  foo = MyModel.objects.get(id=1) # This is `Any` but it should be `MyModel`
-  return foo.xyz # No error, but there should be
+def use_my_model() -> int:
+  foo = MyModel.objects.get(id=1) # Should now be `MyModel`
+  return foo.xyz # Gives an error
 ```
 
-There is a workaround: use `Manager.from_queryset` instead.
+Or if you're using `Manager.from_queryset`:
 
 Example:
 
@@ -177,11 +175,14 @@ class MyModel(models.Model):
   bar = models.IntegerField()
   objects = MyModelManager()
 
-def use_my_model():
-  foo = MyModel.objects.get(id=1)
+def use_my_model() -> int:
+  foo = MyModel.objects.get(id=1) # Should now be `MyModel`
   return foo.xyz # Gives an error
 ```
 
+Take note that `.from_queryset` needs to be placed at _module level_ to annotate
+types correctly. Calling `.from_queryset` in class definition will yield an error.
+
 ### How do I annotate cases where I called QuerySet.annotate?
 
 Django-stubs provides a special type, `django_stubs_ext.WithAnnotations[Model]`, which indicates that the `Model` has

mypy_django_plugin/lib/helpers.py

@@ -376,33 +376,31 @@ def bind_or_analyze_type(t: MypyType, api: SemanticAnalyzer, module_name: Option
 
 
 def copy_method_to_another_class(
-    ctx: ClassDefContext,
+    api: SemanticAnalyzer,
+    cls: ClassDef,
     self_type: Instance,
     new_method_name: str,
     method_node: FuncDef,
     return_type: Optional[MypyType] = None,
     original_module_name: Optional[str] = None,
 ) -> None:
-    semanal_api = get_semanal_api(ctx)
     if method_node.type is None:
-        if not semanal_api.final_iteration:
-            semanal_api.defer()
+        if not api.final_iteration:
+            api.defer()
             return
 
         arguments, return_type = build_unannotated_method_args(method_node)
-        add_method_to_class(
-            semanal_api, ctx.cls, new_method_name, args=arguments, return_type=return_type, self_type=self_type
-        )
+        add_method_to_class(api, cls, new_method_name, args=arguments, return_type=return_type, self_type=self_type)
         return
 
     method_type = method_node.type
     if not isinstance(method_type, CallableType):
-        if not semanal_api.final_iteration:
-            semanal_api.defer()
+        if not api.final_iteration:
+            api.defer()
         return
 
     if return_type is None:
-        return_type = bind_or_analyze_type(method_type.ret_type, semanal_api, original_module_name)
+        return_type = bind_or_analyze_type(method_type.ret_type, api, original_module_name)
     if return_type is None:
         return
 
@@ -415,7 +413,7 @@ def copy_method_to_another_class(
         zip(method_type.arg_types[1:], method_type.arg_kinds[1:], method_type.arg_names[1:]),
         start=1,
     ):
-        bound_arg_type = bind_or_analyze_type(arg_type, semanal_api, original_module_name)
+        bound_arg_type = bind_or_analyze_type(arg_type, api, original_module_name)
         if bound_arg_type is None:
             return
         if arg_name is None and hasattr(method_node, "arguments"):
@@ -431,9 +429,7 @@ def copy_method_to_another_class(
             )
         )
 
-    add_method_to_class(
-        semanal_api, ctx.cls, new_method_name, args=arguments, return_type=return_type, self_type=self_type
-    )
+    add_method_to_class(api, cls, new_method_name, args=arguments, return_type=return_type, self_type=self_type)
 
 
 def add_new_manager_base(api: SemanticAnalyzerPluginInterface, fullname: str) -> None:

mypy_django_plugin/main.py

@@ -23,6 +23,7 @@
 from mypy_django_plugin.lib import fullnames, helpers
 from mypy_django_plugin.transformers import fields, forms, init_create, meta, querysets, request, settings
 from mypy_django_plugin.transformers.managers import (
+    create_new_manager_class_from_as_manager_method,
     create_new_manager_class_from_from_queryset_method,
     fail_if_manager_type_created_in_model_body,
     resolve_manager_method,
@@ -303,11 +304,16 @@ def get_type_analyze_hook(self, fullname: str) -> Optional[Callable[[AnalyzeType
 
     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"):
+        if fullname.endswith(".from_queryset"):
             class_name, _, _ = fullname.rpartition(".")
             info = self._get_typeinfo_or_none(class_name)
             if info and info.has_base(fullnames.BASE_MANAGER_CLASS_FULLNAME):
                 return create_new_manager_class_from_from_queryset_method
+        elif fullname.endswith(".as_manager"):
+            class_name, _, _ = fullname.rpartition(".")
+            info = self._get_typeinfo_or_none(class_name)
+            if info and info.has_base(fullnames.QUERYSET_CLASS_FULLNAME):
+                return create_new_manager_class_from_as_manager_method
         return None