Add common plugin API available for all hooks (python#6044)

ilevkivskyi · web-flow · commit a26b4dfaeadc · 2018-12-13T22:33:04.000Z
There is a problem with `get_base_class_hook()`: only the full name of the base class is passed to it. This causes problems when this hook needs to interact with other hooks. In normal mode this can be solved by using a global plugin state, but this doesn't work in incremental mode, since this plugin state is not stored in cache. There is however already a field that is stored in cache between incremental runs and is free to use by plugins: `TypeInfo.metadata`.

The problem however is that plugin hooks get only full names and can't access the corresponding `TypeInfo`. This is a common enough problem that we should add an API to make the nodes referred by full names accessible to the plugin hooks. This is done by adding a common lookup method to the `Plugin` class.
diff --git a/mypy/build.py b/mypy/build.py
@@ -2386,6 +2386,7 @@ def load_graph(sources: List[BuildSource], manager: BuildManager,
                 if dep in st.suppressed:
                     st.suppressed.remove(dep)
                     st.dependencies.append(dep)
+    manager.plugin.set_modules(manager.modules)
     return graph
 
 
diff --git a/mypy/fixup.py b/mypy/fixup.py
@@ -13,6 +13,7 @@
     TypeType, NOT_READY
 )
 from mypy.visitor import NodeVisitor
+from mypy.lookup import lookup_fully_qualified
 
 
 # N.B: we do a quick_and_dirty fixup in both quick_and_dirty mode and
@@ -257,40 +258,7 @@ def lookup_qualified(modules: Dict[str, MypyFile], name: str,
 
 def lookup_qualified_stnode(modules: Dict[str, MypyFile], name: str,
                             quick_and_dirty: bool) -> Optional[SymbolTableNode]:
-    head = name
-    rest = []
-    while True:
-        if '.' not in head:
-            if not quick_and_dirty:
-                assert '.' in head, "Cannot find %s" % (name,)
-            return None
-        head, tail = head.rsplit('.', 1)
-        rest.append(tail)
-        mod = modules.get(head)
-        if mod is not None:
-            break
-    names = mod.names
-    while True:
-        if not rest:
-            if not quick_and_dirty:
-                assert rest, "Cannot find %s" % (name,)
-            return None
-        key = rest.pop()
-        if key not in names:
-            if not quick_and_dirty:
-                assert key in names, "Cannot find %s for %s" % (key, name)
-            return None
-        stnode = names[key]
-        if not rest:
-            return stnode
-        node = stnode.node
-        # In fine-grained mode, could be a cross-reference to a deleted module
-        # or a Var made up for a missing module.
-        if not isinstance(node, TypeInfo):
-            if not quick_and_dirty:
-                assert node, "Cannot find %s" % (name,)
-            return None
-        names = node.names
+    return lookup_fully_qualified(name, modules, raise_on_missing=not quick_and_dirty)
 
 
 def stale_info(modules: Dict[str, MypyFile]) -> TypeInfo:
diff --git a/mypy/interpreted_plugin.py b/mypy/interpreted_plugin.py
@@ -1,9 +1,10 @@
 """Hack for handling non-mypyc compiled plugins with a mypyc-compiled mypy"""
 
-from typing import Optional, Callable, Any
+from typing import Optional, Callable, Any, Dict
 from mypy.options import Options
 from mypy.types import Type, CallableType
-
+from mypy.nodes import SymbolTableNode, MypyFile
+from mypy.lookup import lookup_fully_qualified
 
 MYPY = False
 if MYPY:
@@ -29,6 +30,14 @@ def __new__(cls, *args: Any, **kwargs: Any) -> 'mypy.plugin.Plugin':
     def __init__(self, options: Options) -> None:
         self.options = options
         self.python_version = options.python_version
+        self._modules = None  # type: Optional[Dict[str, MypyFile]]
+
+    def set_modules(self, modules: Dict[str, MypyFile]) -> None:
+        self._modules = modules
+
+    def lookup_fully_qualified(self, fullname: str) -> Optional[SymbolTableNode]:
+        assert self._modules is not None
+        return lookup_fully_qualified(fullname, self._modules)
 
     def get_type_analyze_hook(self, fullname: str
                               ) -> Optional[Callable[['mypy.plugin.AnalyzeTypeContext'], Type]]:
diff --git a/mypy/lookup.py b/mypy/lookup.py
@@ -0,0 +1,59 @@
+"""
+This is a module for various lookup functions:
+functions that will find a semantic node by its name.
+"""
+
+from mypy.nodes import MypyFile, SymbolTableNode, TypeInfo
+from typing import Dict, Optional
+
+# TODO: gradually move existing lookup functions to this module.
+
+
+def lookup_fully_qualified(name: str, modules: Dict[str, MypyFile],
+                           raise_on_missing: bool = False) -> Optional[SymbolTableNode]:
+    """Find a symbol using it fully qualified name.
+
+    The algorithm has two steps: first we try splitting the name on '.' to find
+    the module, then iteratively look for each next chunk after a '.' (e.g. for
+    nested classes).
+
+    This function should *not* be used to find a module. Those should be looked
+    in the modules dictionary.
+    """
+    head = name
+    rest = []
+    # 1. Find a module tree in modules dictionary.
+    while True:
+        if '.' not in head:
+            if raise_on_missing:
+                assert '.' in head, "Cannot find module for %s" % (name,)
+            return None
+        head, tail = head.rsplit('.', maxsplit=1)
+        rest.append(tail)
+        mod = modules.get(head)
+        if mod is not None:
+            break
+    names = mod.names
+    # 2. Find the symbol in the module tree.
+    if not rest:
+        # Looks like a module, don't use this to avoid confusions.
+        if raise_on_missing:
+            assert rest, "Cannot find %s, got a module symbol" % (name,)
+        return None
+    while True:
+        key = rest.pop()
+        if key not in names:
+            if raise_on_missing:
+                assert key in names, "Cannot find %s for %s" % (key, name)
+            return None
+        stnode = names[key]
+        if not rest:
+            return stnode
+        node = stnode.node
+        # In fine-grained mode, could be a cross-reference to a deleted module
+        # or a Var made up for a missing module.
+        if not isinstance(node, TypeInfo):
+            if raise_on_missing:
+                assert node, "Cannot find %s" % (name,)
+            return None
+        names = node.names
diff --git a/mypy/plugin.py b/mypy/plugin.py
@@ -43,11 +43,14 @@
 from typing import Callable, List, Tuple, Optional, NamedTuple, TypeVar, Dict
 from mypy_extensions import trait
 
-from mypy.nodes import Expression, Context, ClassDef, SymbolTableNode, MypyFile, CallExpr
+from mypy.nodes import (
+    Expression, Context, ClassDef, SymbolTableNode, MypyFile, CallExpr
+)
 from mypy.tvar_scope import TypeVarScope
 from mypy.types import Type, Instance, CallableType, TypeList, UnboundType
 from mypy.messages import MessageBuilder
 from mypy.options import Options
+from mypy.lookup import lookup_fully_qualified
 import mypy.interpreted_plugin
 
 
@@ -96,6 +99,28 @@ def analyze_callable_args(self, arglist: TypeList) -> Optional[Tuple[List[Type],
         ('api', TypeAnalyzerPluginInterface)])
 
 
+@trait
+class CommonPluginApi:
+    """
+    A common plugin API (shared between semantic analysis and type checking phases)
+    that all plugin hooks get independently of the context.
+    """
+
+    # Global mypy options.
+    # Per-file options can be only accessed on various
+    # XxxPluginInterface classes.
+    options = None  # type: Options
+
+    @abstractmethod
+    def lookup_fully_qualified(self, fullname: str) -> Optional[SymbolTableNode]:
+        """Lookup a symbol by its full name (including module).
+
+        This lookup function available for all plugins. Return None if a name
+        is not found. This function doesn't support lookup from current scope.
+        Use SemanticAnalyzerPluginInterface.lookup_qualified() for this."""
+        raise NotImplementedError
+
+
 @trait
 class CheckerPluginInterface:
     """Interface for accessing type checker functionality in plugins.
@@ -286,7 +311,7 @@ def qualified_name(self, n: str) -> str:
     ])
 
 
-class Plugin:
+class Plugin(CommonPluginApi):
     """Base class of all type checker plugins.
 
     This defines a no-op plugin.  Subclasses can override some methods to
@@ -303,6 +328,17 @@ class Plugin:
     def __init__(self, options: Options) -> None:
         self.options = options
         self.python_version = options.python_version
+        # This can't be set in __init__ because it is executed too soon in build.py.
+        # Therefore, build.py *must* set it later before graph processing starts
+        # by calling set_modules().
+        self._modules = None  # type: Optional[Dict[str, MypyFile]]
+
+    def set_modules(self, modules: Dict[str, MypyFile]) -> None:
+        self._modules = modules
+
+    def lookup_fully_qualified(self, fullname: str) -> Optional[SymbolTableNode]:
+        assert self._modules is not None
+        return lookup_fully_qualified(fullname, self._modules)
 
     def get_type_analyze_hook(self, fullname: str
                               ) -> Optional[Callable[[AnalyzeTypeContext], Type]]:
@@ -479,6 +515,12 @@ def __init__(self, plugin: mypy.interpreted_plugin.InterpretedPlugin) -> None:
         super().__init__(plugin.options)
         self.plugin = plugin
 
+    def set_modules(self, modules: Dict[str, MypyFile]) -> None:
+        self.plugin.set_modules(modules)
+
+    def lookup_fully_qualified(self, fullname: str) -> Optional[SymbolTableNode]:
+        return self.plugin.lookup_fully_qualified(fullname)
+
     def get_type_analyze_hook(self, fullname: str
                               ) -> Optional[Callable[[AnalyzeTypeContext], Type]]:
         return self.plugin.get_type_analyze_hook(fullname)
@@ -540,6 +582,10 @@ def __init__(self, options: Options, plugins: List[Plugin]) -> None:
         super().__init__(options)
         self._plugins = plugins
 
+    def set_modules(self, modules: Dict[str, MypyFile]) -> None:
+        for plugin in self._plugins:
+            plugin.set_modules(modules)
+
     def get_type_analyze_hook(self, fullname: str
                               ) -> Optional[Callable[[AnalyzeTypeContext], Type]]:
         return self._find_hook(lambda plugin: plugin.get_type_analyze_hook(fullname))
diff --git a/mypy/semanal.py b/mypy/semanal.py
@@ -3504,6 +3504,9 @@ def lookup_fully_qualified_or_none(self, fullname: str) -> Optional[SymbolTableN
 
         Note that this can't be used for names nested in class namespaces.
         """
+        # TODO: unify/clean-up/simplify lookup methods, see #4157.
+        # TODO: support nested classes (but consider performance impact,
+        #       we might keep the module level only lookup for thing like 'builtins.int').
         assert '.' in fullname
         module, name = fullname.rsplit('.', maxsplit=1)
         if module not in self.modules:
diff --git a/test-data/unit/check-custom-plugin.test b/test-data/unit/check-custom-plugin.test
@@ -331,3 +331,34 @@ class Instr(Generic[T]): ...
 [file mypy.ini]
 [[mypy]
 plugins=<ROOT>/test-data/unit/plugins/dyn_class.py
+
+[case testBaseClassPluginHookWorksIncremental]
+# flags: --config-file tmp/mypy.ini
+import a
+
+[file a.py]
+from base import Base
+class C(Base): ...
+
+[file a.py.2]
+from base import Base
+class C(Base): ...
+reveal_type(C().__magic__)
+Base.__magic__
+
+[file base.py]
+from lib import declarative_base
+Base = declarative_base()
+
+[file lib.py]
+from typing import Any
+def declarative_base() -> Any: ...
+
+[file mypy.ini]
+[[mypy]
+python_version=3.6
+plugins=<ROOT>/test-data/unit/plugins/common_api_incremental.py
+[out]
+[out2]
+tmp/a.py:3: error: Revealed type is 'builtins.str'
+tmp/a.py:4: error: "Type[Base]" has no attribute "__magic__"
diff --git a/test-data/unit/plugins/common_api_incremental.py b/test-data/unit/plugins/common_api_incremental.py
@@ -0,0 +1,44 @@
+from mypy.plugin import Plugin
+from mypy.nodes import (
+    ClassDef, Block, TypeInfo, SymbolTable, SymbolTableNode, MDEF, GDEF, Var
+)
+
+
+class DynPlugin(Plugin):
+    def get_dynamic_class_hook(self, fullname):
+        if fullname == 'lib.declarative_base':
+            return add_info_hook
+        return None
+
+    def get_base_class_hook(self, fullname: str):
+        sym = self.lookup_fully_qualified(fullname)
+        if sym and isinstance(sym.node, TypeInfo):
+            if sym.node.metadata.get('magic'):
+                return add_magic_hook
+        return None
+
+
+def add_info_hook(ctx) -> 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
+    obj = ctx.api.builtin_type('builtins.object')
+    info.mro = [info, obj.type]
+    info.bases = [obj]
+    ctx.api.add_symbol_table_node(ctx.name, SymbolTableNode(GDEF, info))
+    info.metadata['magic'] = True
+
+
+def add_magic_hook(ctx) -> None:
+    info = ctx.cls.info
+    str_type = ctx.api.named_type_or_none('builtins.str', [])
+    assert str_type is not None
+    var = Var('__magic__', str_type)
+    var.info = info
+    info.names['__magic__'] = SymbolTableNode(MDEF, var)
+
+
+def plugin(version):
+    return DynPlugin
