diff --git a/mypy/semanal.py b/mypy/semanal.py index c18b44394002..718da41937af 100644 --- a/mypy/semanal.py +++ b/mypy/semanal.py @@ -2194,9 +2194,15 @@ def apply_dynamic_class_hook(self, s: AssignmentStmt) -> None: if not isinstance(lval, NameExpr) or not isinstance(s.rvalue, CallExpr): return call = s.rvalue - if not isinstance(call.callee, RefExpr): - return - fname = call.callee.fullname + fname = None + if isinstance(call.callee, RefExpr): + fname = call.callee.fullname + # check if method call + if fname is None and isinstance(call.callee, MemberExpr): + callee_expr = call.callee.expr + if isinstance(callee_expr, RefExpr) and callee_expr.fullname: + method_name = call.callee.name + fname = callee_expr.fullname + '.' + method_name if fname: hook = self.plugin.get_dynamic_class_hook(fname) if hook: diff --git a/test-data/unit/check-custom-plugin.test b/test-data/unit/check-custom-plugin.test index 0e9d517b41bb..d2e2221ef5e3 100644 --- a/test-data/unit/check-custom-plugin.test +++ b/test-data/unit/check-custom-plugin.test @@ -516,6 +516,35 @@ class Instr(Generic[T]): ... \[mypy] plugins=/test-data/unit/plugins/dyn_class.py +[case testDynamicClassHookFromClassMethod] +# flags: --config-file tmp/mypy.ini + +from mod import QuerySet, Manager + +MyManager = Manager.from_queryset(QuerySet) + +reveal_type(MyManager()) # N: Revealed type is '__main__.MyManager' +reveal_type(MyManager().attr) # N: Revealed type is 'builtins.str' + +def func(manager: MyManager) -> None: + reveal_type(manager) # N: Revealed type is '__main__.MyManager' + reveal_type(manager.attr) # N: Revealed type is 'builtins.str' + +func(MyManager()) + +[file mod.py] +from typing import Generic, TypeVar, Type +class QuerySet: + attr: str +class Manager: + @classmethod + def from_queryset(cls, queryset_cls: Type[QuerySet]): ... + +[builtins fixtures/classmethod.pyi] +[file mypy.ini] +\[mypy] +plugins=/test-data/unit/plugins/dyn_class_from_method.py + [case testBaseClassPluginHookWorksIncremental] # flags: --config-file tmp/mypy.ini import a diff --git a/test-data/unit/plugins/dyn_class_from_method.py b/test-data/unit/plugins/dyn_class_from_method.py new file mode 100644 index 000000000000..8a18f7f1e8e1 --- /dev/null +++ b/test-data/unit/plugins/dyn_class_from_method.py @@ -0,0 +1,28 @@ +from mypy.nodes import (Block, ClassDef, GDEF, SymbolTable, SymbolTableNode, TypeInfo) +from mypy.plugin import DynamicClassDefContext, Plugin +from mypy.types import Instance + + +class DynPlugin(Plugin): + def get_dynamic_class_hook(self, fullname): + if 'from_queryset' in fullname: + return add_info_hook + return None + + +def add_info_hook(ctx: DynamicClassDefContext): + 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 + queryset_type_fullname = ctx.call.args[0].fullname + queryset_info = ctx.api.lookup_fully_qualified_or_none(queryset_type_fullname).node # type: TypeInfo + obj = ctx.api.builtin_type('builtins.object') + info.mro = [info, queryset_info, obj.type] + info.bases = [Instance(queryset_info, [])] + ctx.api.add_symbol_table_node(ctx.name, SymbolTableNode(GDEF, info)) + + +def plugin(version): + return DynPlugin