[WIP] Add get_function_decorator_plugin_hook to plugin interface

mypy/plugin.py

@@ -124,7 +124,8 @@ class C: pass
 from mypy_extensions import trait, mypyc_attr
 
 from mypy.nodes import (
-    Expression, Context, ClassDef, SymbolTableNode, MypyFile, CallExpr, ArgKind, TypeInfo
+    Expression, Context, ClassDef, SymbolTableNode, MypyFile, CallExpr, ArgKind, TypeInfo,
+    Decorator
 )
 from mypy.tvar_scope import TypeVarLikeScope
 from mypy.types import (
@@ -475,6 +476,16 @@ class DynamicClassDefContext(NamedTuple):
     api: SemanticAnalyzerPluginInterface
 
 
+# A context for a decorator hook, that modifies the function definition
+FunctionDecoratorContext = NamedTuple(
+    'FunctionDecoratorContext', [
+        ('decorator', Expression),
+        ('decorated_function', Decorator),
+        ('api', SemanticAnalyzerPluginInterface)
+    ]
+)
+
+
 @mypyc_attr(allow_interpreted_subclasses=True)
 class Plugin(CommonPluginApi):
     """Base class of all type checker plugins.
@@ -745,6 +756,18 @@ def get_dynamic_class_hook(self, fullname: str
         """
         return None
 
+    def get_function_decorator_hook(self, fullname: str
+                                    ) -> Optional[Callable[[FunctionDecoratorContext], bool]]:
+        """Update function definition for given function decorators
+
+        The plugin can modify a function _in place_.
+
+        The hook is called with full names of all function decorators.
+
+        Return true if the decorator has been handled and should be removed
+        """
+        return None
+
 
 T = TypeVar('T')
 
@@ -831,6 +854,10 @@ def get_dynamic_class_hook(self, fullname: str
                                ) -> Optional[Callable[[DynamicClassDefContext], None]]:
         return self._find_hook(lambda plugin: plugin.get_dynamic_class_hook(fullname))
 
+    def get_function_decorator_hook(self, fullname: str
+                                    ) -> Optional[Callable[[FunctionDecoratorContext], bool]]:
+        return self._find_hook(lambda plugin: plugin.get_function_decorator_hook(fullname))
+
     def _find_hook(self, lookup: Callable[[Plugin], T]) -> Optional[T]:
         for plugin in self._plugins:
             hook = lookup(plugin)

mypy/semanal.py

@@ -112,7 +112,7 @@
 from mypy.options import Options
 from mypy.plugin import (
     Plugin, ClassDefContext, SemanticAnalyzerPluginInterface,
-    DynamicClassDefContext
+    DynamicClassDefContext, FunctionDecoratorContext
 )
 from mypy.util import (
     correct_relative_import, unmangle, module_prefix, is_typeshed_file, unnamed_function,
@@ -1094,6 +1094,8 @@ def visit_decorator(self, dec: Decorator) -> None:
                     removed.append(i)
                 else:
                     self.fail("@final cannot be used with non-method functions", d)
+            if self.apply_decorator_plugin_hooks(d, dec):
+                removed.append(i)
         for i in reversed(removed):
             del dec.decorators[i]
         if (not dec.is_overload or dec.var.is_property) and self.type:
@@ -1111,6 +1113,31 @@ def check_decorated_function_is_method(self, decorator: str,
         if not self.type or self.is_func_scope():
             self.fail(f'"{decorator}" used with a non-method', context)
 
+    def apply_decorator_plugin_hooks(self, node: Expression, dec: Decorator) -> bool:
+        # TODO: Remove duplicate code
+        def get_fullname(expr: Expression) -> Optional[str]:
+            if isinstance(expr, CallExpr):
+                return get_fullname(expr.callee)
+            elif isinstance(expr, IndexExpr):
+                return get_fullname(expr.base)
+            elif isinstance(expr, RefExpr):
+                if expr.fullname:
+                    return expr.fullname
+                # If we don't have a fullname look it up. This happens because base classes are
+                # analyzed in a different manner (see exprtotype.py) and therefore those AST
+                # nodes will not have full names.
+                sym = self.lookup_type_node(expr)
+                if sym:
+                    return sym.fullname
+            return None
+
+        decorator_name = get_fullname(node)
+        if decorator_name:
+            hook = self.plugin.get_function_decorator_hook(decorator_name)
+            if hook:
+                return hook(FunctionDecoratorContext(node, dec, self))
+        return False
+
     #
     # Classes
     #

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

@@ -991,3 +991,40 @@ class Cls:
 [file mypy.ini]
 \[mypy]
 plugins=<ROOT>/test-data/unit/plugins/class_attr_hook.py
+
+[case testFunctionDecoratorPluginHookForFunction]
+# flags: --config-file tmp/mypy.ini
+
+from m import decorator
+
+@decorator
+def function(self) -> str: ...
+
+@function.setter
+def function(self, value: str) -> None: ...
+
+[file m.py]
+from typing import Callable
+def decorator(param) -> Callable[..., str]: pass
+[file mypy.ini]
+\[mypy]
+plugins=<ROOT>/test-data/unit/plugins/function_decorator_hook.py
+
+[case testFunctionDecoratorPluginHookForMethod]
+# flags: --config-file tmp/mypy.ini
+
+from m import decorator
+
+class A:
+    @decorator
+    def property(self) -> str: ...
+
+    @property.setter
+    def property(self, value: str) -> None: ...
+
+[file m.py]
+from typing import Callable
+def decorator(param) -> Callable[..., str]: pass
+[file mypy.ini]
+\[mypy]
+plugins=<ROOT>/test-data/unit/plugins/function_decorator_hook.py

test-data/unit/plugins/function_decorator_hook.py

@@ -0,0 +1,19 @@
+from mypy.plugin import Plugin, FunctionDecoratorContext
+
+
+class FunctionDecoratorPlugin(Plugin):
+    def get_function_decorator_hook(self, fullname):
+        if fullname == 'm.decorator':
+            return my_hook
+        return None
+
+
+def my_hook(ctx: FunctionDecoratorContext) -> bool:
+    ctx.decorated_function.func.is_property = True
+    ctx.decorated_function.var.is_property = True
+
+    return True
+
+
+def plugin(version):
+    return FunctionDecoratorPlugin