[WIP] Add get_function_decorator_plugin_hook to plugin interface

mypy/plugin.py

@@ -124,7 +124,7 @@ class C: pass
 from mypy_extensions import trait, mypyc_attr
 
 from mypy.nodes import (
- Expression, Context, ClassDef, SymbolTableNode, MypyFile, CallExpr
+ Expression, Context, ClassDef, SymbolTableNode, MypyFile, CallExpr, Decorator
 )
 from mypy.tvar_scope import TypeVarLikeScope
 from mypy.types import Type, Instance, CallableType, TypeList, UnboundType, ProperType
@@ -454,6 +454,16 @@ def final_iteration(self) -> bool:
  ])
 
 
+# A context for a decorator hook, that modifies the function definition
+FunctionDecoratorContext = NamedTuple(
+ 'FunctionDecoratorContext', [
+ ('decorator', Expression),
+ ('decoratedFunction', Decorator),
+ ('api', SemanticAnalyzerPluginInterface)
+ ]
+)
+
+
 @mypyc_attr(allow_interpreted_subclasses=True)
 class Plugin(CommonPluginApi):
  """Base class of all type checker plugins.
@@ -705,6 +715,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')
 
@@ -787,6 +809,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

@@ -104,7 +104,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
 from mypy.scope import Scope
@@ -1023,6 +1023,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:
@@ -1038,6 +1040,31 @@ def check_decorated_function_is_method(self, decorator: str,
  if not self.type or self.is_func_scope():
  self.fail("'%s' used with a non-method" % decorator, 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

@@ -730,3 +730,40 @@ reveal_type(dynamic_signature(1)) # N: Revealed type is 'builtins.int'
 [file mypy.ini]
 \[mypy]
 plugins=<ROOT>/test-data/unit/plugins/function_sig_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.decoratedFunction.func.is_property = True
+ ctx.decoratedFunction.var.is_property = True
+
+ return True
+
+
+def plugin(version):
+ return FunctionDecoratorPlugin