Skip to content

Commit

Permalink
Invoke get_dynamic_class_hook on method calls (#7990)
Browse files Browse the repository at this point in the history
Fixes #7266

Note: a RefExpr (including MemberExpr) with a non-None full name should always take precedence.
  • Loading branch information
syastrov authored and ilevkivskyi committed Nov 23, 2019
1 parent fffb4b4 commit 20b60b5
Show file tree
Hide file tree
Showing 3 changed files with 66 additions and 3 deletions.
12 changes: 9 additions & 3 deletions mypy/semanal.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
29 changes: 29 additions & 0 deletions test-data/unit/check-custom-plugin.test
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
28 changes: 28 additions & 0 deletions test-data/unit/plugins/dyn_class_from_method.py
Original file line number Diff line number Diff line change
@@ -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

0 comments on commit 20b60b5

Please sign in to comment.