fixup! fixup! fixup! fixup! Implement support for <QuerySet>.as_manager()

Author: flaeppe

mypy_django_plugin/main.py

@@ -304,13 +304,12 @@ 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"):
-            class_name, _, _ = fullname.rpartition(".")
+        class_name, _, method_name = fullname.rpartition(".")
+        if method_name == "from_queryset":
             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(".")
+        elif method_name == "as_manager":
             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

mypy_django_plugin/transformers/managers.py

@@ -286,13 +286,15 @@ def create_new_manager_class_from_from_queryset_method(ctx: DynamicClassDefConte
     queryset_info = queryset_sym.node
     assert isinstance(queryset_info, TypeInfo)
 
+    class_name_param_value = None
     if len(ctx.call.args) > 1:
         expr = ctx.call.args[1]
-        assert isinstance(expr, StrExpr)
-        manager_class_name = expr.value
-    else:
-        manager_class_name = manager_base.name + "From" + queryset_info.name
+        if isinstance(expr, StrExpr):
+            class_name_param_value = expr.value
+        # When the argument is not a `StrExpr` we're gonna have a hard time figuring
+        # out what the runtime value will be, so we'll just use the default instead.
 
+    manager_class_name = class_name_param_value or manager_base.name + "From" + queryset_info.name
     manager_base_instance = fill_typevars(manager_base)
     assert isinstance(manager_base_instance, Instance)
     # Create a new `TypeInfo` instance for the manager type
@@ -324,8 +326,8 @@ def create_new_manager_class_from_from_queryset_method(ctx: DynamicClassDefConte
     var._fullname = f"{semanal_api.cur_mod_id}.{ctx.name}"
     var.is_inferred = True
     # Note: Order of `add_symbol_table_node` calls matter. Case being if
-    # `ctx.name == new_manager_info.name`, then we'd _only_ like the type and not the
-    # `Var` to exist..
+    # `ctx.name == new_manager_info.name`, then we'd like the type to override the
+    # `Var`, so the `Var` won't exist in the end.
     assert semanal_api.add_symbol_table_node(ctx.name, SymbolTableNode(symbol_kind, var, plugin_generated=True))
     # Insert the new manager dynamic class
     assert semanal_api.add_symbol_table_node(

tests/typecheck/managers/querysets/test_from_queryset.yml

@@ -545,3 +545,41 @@
               BaseManagerFromModelQuerySet = BaseManager.from_queryset(ModelQuerySet)
               class MyModel(models.Model):
                   objects = BaseManagerFromModelQuerySet()
+
+-   case: accepts_explicit_none_as_class_name
+    main: |
+        from myapp.models import PositionalNone, NoneAsKwarg
+        reveal_type(PositionalNone)  # N: Revealed type is "Type[myapp.models.BaseManagerFromModelQuerySet[Any]]"
+        reveal_type(NoneAsKwarg)  # N: Revealed type is "Type[myapp.models.BaseManagerFromModelQuerySet[Any]]"
+    installed_apps:
+        - myapp
+    files:
+        - path: myapp/__init__.py
+        - path: myapp/models.py
+          content: |
+              from django.db import models
+              from django.db.models.manager import BaseManager
+
+              class ModelQuerySet(models.QuerySet):
+                  ...
+
+              PositionalNone = BaseManager.from_queryset(ModelQuerySet, None)
+              NoneAsKwarg = BaseManager.from_queryset(ModelQuerySet, class_name=None)
+
+-   case: uses_fallback_class_name_when_argument_is_not_string_expression
+    main: |
+        from myapp.models import StrCallable
+        reveal_type(StrCallable)  # N: Revealed type is "Type[myapp.models.BaseManagerFromModelQuerySet[Any]]"
+    installed_apps:
+        - myapp
+    files:
+        - path: myapp/__init__.py
+        - path: myapp/models.py
+          content: |
+              from django.db import models
+              from django.db.models.manager import BaseManager
+
+              class ModelQuerySet(models.QuerySet):
+                  ...
+
+              StrCallable = BaseManager.from_queryset(ModelQuerySet, class_name=str(1))