Invoke get_dynamic_class_hook on method calls (python#7990)

syastrov · ilevkivskyi · commit 20b60b567889 · 2019-11-23T23:24:43.000Z
Fixes python#7266 Note: a RefExpr (including MemberExpr) with a non-None full name should always take precedence.
diff --git 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
@@ -516,6 +516,35 @@ class Instr(Generic[T]): ...
 \[mypy]
 plugins=<ROOT>/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=<ROOT>/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
@@ -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
