From ba978f461e1f88327f9caa2e83774caaaeee1a6a Mon Sep 17 00:00:00 2001
From: Petter Friberg <petter@5monkeys.se>
Date: Tue, 19 Sep 2023 08:10:31 +0200
Subject: [PATCH] Call dynamic class hook on generic classes (#16052)

Fixes: #8359

CC @sobolevn

`get_dynamic_class_hook()` will now additionally be called for generic
classes with parameters. e.g.

```python
y = SomeGenericClass[type, ...].method()
```
---
 mypy/semanal.py                               |  7 ++++
 test-data/unit/check-custom-plugin.test       | 12 +++++-
 .../unit/plugins/dyn_class_from_method.py     | 40 ++++++++++++++++++-
 3 files changed, 57 insertions(+), 2 deletions(-)

diff --git a/mypy/semanal.py b/mypy/semanal.py
index 70403eed57ae..e19cd86d5e89 100644
--- a/mypy/semanal.py
+++ b/mypy/semanal.py
@@ -3205,6 +3205,13 @@ def apply_dynamic_class_hook(self, s: AssignmentStmt) -> None:
                 if isinstance(callee_expr, RefExpr) and callee_expr.fullname:
                     method_name = call.callee.name
                     fname = callee_expr.fullname + "." + method_name
+                elif (
+                    isinstance(callee_expr, IndexExpr)
+                    and isinstance(callee_expr.base, RefExpr)
+                    and isinstance(callee_expr.analyzed, TypeApplication)
+                ):
+                    method_name = call.callee.name
+                    fname = callee_expr.base.fullname + "." + method_name
                 elif isinstance(callee_expr, CallExpr):
                     # check if chain call
                     call = callee_expr
diff --git a/test-data/unit/check-custom-plugin.test b/test-data/unit/check-custom-plugin.test
index 22374d09cf9f..63529cf165ce 100644
--- a/test-data/unit/check-custom-plugin.test
+++ b/test-data/unit/check-custom-plugin.test
@@ -684,12 +684,16 @@ plugins=<ROOT>/test-data/unit/plugins/dyn_class.py
 [case testDynamicClassHookFromClassMethod]
 # flags: --config-file tmp/mypy.ini
 
-from mod import QuerySet, Manager
+from mod import QuerySet, Manager, GenericQuerySet
 
 MyManager = Manager.from_queryset(QuerySet)
+ManagerFromGenericQuerySet = GenericQuerySet[int].as_manager()
 
 reveal_type(MyManager())  # N: Revealed type is "__main__.MyManager"
 reveal_type(MyManager().attr)  # N: Revealed type is "builtins.str"
+reveal_type(ManagerFromGenericQuerySet())  # N: Revealed type is "__main__.ManagerFromGenericQuerySet"
+reveal_type(ManagerFromGenericQuerySet().attr)  # N: Revealed type is "builtins.int"
+queryset: GenericQuerySet[int] = ManagerFromGenericQuerySet()
 
 def func(manager: MyManager) -> None:
     reveal_type(manager)   # N: Revealed type is "__main__.MyManager"
@@ -704,6 +708,12 @@ class QuerySet:
 class Manager:
     @classmethod
     def from_queryset(cls, queryset_cls: Type[QuerySet]): ...
+T = TypeVar("T")
+class GenericQuerySet(Generic[T]):
+    attr: T
+
+    @classmethod
+    def as_manager(cls): ...
 
 [builtins fixtures/classmethod.pyi]
 [file mypy.ini]
diff --git a/test-data/unit/plugins/dyn_class_from_method.py b/test-data/unit/plugins/dyn_class_from_method.py
index b84754654084..2630b16be66e 100644
--- a/test-data/unit/plugins/dyn_class_from_method.py
+++ b/test-data/unit/plugins/dyn_class_from_method.py
@@ -2,7 +2,19 @@
 
 from typing import Callable
 
-from mypy.nodes import GDEF, Block, ClassDef, RefExpr, SymbolTable, SymbolTableNode, TypeInfo
+from mypy.nodes import (
+    GDEF,
+    Block,
+    ClassDef,
+    IndexExpr,
+    MemberExpr,
+    NameExpr,
+    RefExpr,
+    SymbolTable,
+    SymbolTableNode,
+    TypeApplication,
+    TypeInfo,
+)
 from mypy.plugin import DynamicClassDefContext, Plugin
 from mypy.types import Instance
 
@@ -13,6 +25,8 @@ def get_dynamic_class_hook(
     ) -> Callable[[DynamicClassDefContext], None] | None:
         if "from_queryset" in fullname:
             return add_info_hook
+        if "as_manager" in fullname:
+            return as_manager_hook
         return None
 
 
@@ -34,5 +48,29 @@ def add_info_hook(ctx: DynamicClassDefContext) -> None:
     ctx.api.add_symbol_table_node(ctx.name, SymbolTableNode(GDEF, info))
 
 
+def as_manager_hook(ctx: DynamicClassDefContext) -> None:
+    class_def = ClassDef(ctx.name, Block([]))
+    class_def.fullname = ctx.api.qualified_name(ctx.name)
+
+    info = TypeInfo(SymbolTable(), class_def, ctx.api.cur_mod_id)
+    class_def.info = info
+    assert isinstance(ctx.call.callee, MemberExpr)
+    assert isinstance(ctx.call.callee.expr, IndexExpr)
+    assert isinstance(ctx.call.callee.expr.analyzed, TypeApplication)
+    assert isinstance(ctx.call.callee.expr.analyzed.expr, NameExpr)
+
+    queryset_type_fullname = ctx.call.callee.expr.analyzed.expr.fullname
+    queryset_node = ctx.api.lookup_fully_qualified_or_none(queryset_type_fullname)
+    assert queryset_node is not None
+    queryset_info = queryset_node.node
+    assert isinstance(queryset_info, TypeInfo)
+    parameter_type = ctx.call.callee.expr.analyzed.types[0]
+
+    obj = ctx.api.named_type("builtins.object")
+    info.mro = [info, queryset_info, obj.type]
+    info.bases = [Instance(queryset_info, [parameter_type])]
+    ctx.api.add_symbol_table_node(ctx.name, SymbolTableNode(GDEF, info))
+
+
 def plugin(version: str) -> type[DynPlugin]:
     return DynPlugin