Skip to content

Commit ba978f4

Browse files
author
Petter Friberg
authored
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() ```
1 parent 1dcff0d commit ba978f4

File tree

3 files changed

+57
-2
lines changed

3 files changed

+57
-2
lines changed

mypy/semanal.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3205,6 +3205,13 @@ def apply_dynamic_class_hook(self, s: AssignmentStmt) -> None:
32053205
if isinstance(callee_expr, RefExpr) and callee_expr.fullname:
32063206
method_name = call.callee.name
32073207
fname = callee_expr.fullname + "." + method_name
3208+
elif (
3209+
isinstance(callee_expr, IndexExpr)
3210+
and isinstance(callee_expr.base, RefExpr)
3211+
and isinstance(callee_expr.analyzed, TypeApplication)
3212+
):
3213+
method_name = call.callee.name
3214+
fname = callee_expr.base.fullname + "." + method_name
32083215
elif isinstance(callee_expr, CallExpr):
32093216
# check if chain call
32103217
call = callee_expr

test-data/unit/check-custom-plugin.test

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -684,12 +684,16 @@ plugins=<ROOT>/test-data/unit/plugins/dyn_class.py
684684
[case testDynamicClassHookFromClassMethod]
685685
# flags: --config-file tmp/mypy.ini
686686

687-
from mod import QuerySet, Manager
687+
from mod import QuerySet, Manager, GenericQuerySet
688688

689689
MyManager = Manager.from_queryset(QuerySet)
690+
ManagerFromGenericQuerySet = GenericQuerySet[int].as_manager()
690691

691692
reveal_type(MyManager()) # N: Revealed type is "__main__.MyManager"
692693
reveal_type(MyManager().attr) # N: Revealed type is "builtins.str"
694+
reveal_type(ManagerFromGenericQuerySet()) # N: Revealed type is "__main__.ManagerFromGenericQuerySet"
695+
reveal_type(ManagerFromGenericQuerySet().attr) # N: Revealed type is "builtins.int"
696+
queryset: GenericQuerySet[int] = ManagerFromGenericQuerySet()
693697

694698
def func(manager: MyManager) -> None:
695699
reveal_type(manager) # N: Revealed type is "__main__.MyManager"
@@ -704,6 +708,12 @@ class QuerySet:
704708
class Manager:
705709
@classmethod
706710
def from_queryset(cls, queryset_cls: Type[QuerySet]): ...
711+
T = TypeVar("T")
712+
class GenericQuerySet(Generic[T]):
713+
attr: T
714+
715+
@classmethod
716+
def as_manager(cls): ...
707717

708718
[builtins fixtures/classmethod.pyi]
709719
[file mypy.ini]

test-data/unit/plugins/dyn_class_from_method.py

Lines changed: 39 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,19 @@
22

33
from typing import Callable
44

5-
from mypy.nodes import GDEF, Block, ClassDef, RefExpr, SymbolTable, SymbolTableNode, TypeInfo
5+
from mypy.nodes import (
6+
GDEF,
7+
Block,
8+
ClassDef,
9+
IndexExpr,
10+
MemberExpr,
11+
NameExpr,
12+
RefExpr,
13+
SymbolTable,
14+
SymbolTableNode,
15+
TypeApplication,
16+
TypeInfo,
17+
)
618
from mypy.plugin import DynamicClassDefContext, Plugin
719
from mypy.types import Instance
820

@@ -13,6 +25,8 @@ def get_dynamic_class_hook(
1325
) -> Callable[[DynamicClassDefContext], None] | None:
1426
if "from_queryset" in fullname:
1527
return add_info_hook
28+
if "as_manager" in fullname:
29+
return as_manager_hook
1630
return None
1731

1832

@@ -34,5 +48,29 @@ def add_info_hook(ctx: DynamicClassDefContext) -> None:
3448
ctx.api.add_symbol_table_node(ctx.name, SymbolTableNode(GDEF, info))
3549

3650

51+
def as_manager_hook(ctx: DynamicClassDefContext) -> None:
52+
class_def = ClassDef(ctx.name, Block([]))
53+
class_def.fullname = ctx.api.qualified_name(ctx.name)
54+
55+
info = TypeInfo(SymbolTable(), class_def, ctx.api.cur_mod_id)
56+
class_def.info = info
57+
assert isinstance(ctx.call.callee, MemberExpr)
58+
assert isinstance(ctx.call.callee.expr, IndexExpr)
59+
assert isinstance(ctx.call.callee.expr.analyzed, TypeApplication)
60+
assert isinstance(ctx.call.callee.expr.analyzed.expr, NameExpr)
61+
62+
queryset_type_fullname = ctx.call.callee.expr.analyzed.expr.fullname
63+
queryset_node = ctx.api.lookup_fully_qualified_or_none(queryset_type_fullname)
64+
assert queryset_node is not None
65+
queryset_info = queryset_node.node
66+
assert isinstance(queryset_info, TypeInfo)
67+
parameter_type = ctx.call.callee.expr.analyzed.types[0]
68+
69+
obj = ctx.api.named_type("builtins.object")
70+
info.mro = [info, queryset_info, obj.type]
71+
info.bases = [Instance(queryset_info, [parameter_type])]
72+
ctx.api.add_symbol_table_node(ctx.name, SymbolTableNode(GDEF, info))
73+
74+
3775
def plugin(version: str) -> type[DynPlugin]:
3876
return DynPlugin

0 commit comments

Comments
 (0)