diff --git a/mypy/plugin.py b/mypy/plugin.py index 72d28c39436a..bf8909778eed 100644 --- a/mypy/plugin.py +++ b/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) diff --git a/mypy/semanal.py b/mypy/semanal.py index 555cb749074e..ffea4ed81fa1 100644 --- a/mypy/semanal.py +++ b/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 # diff --git a/test-data/unit/check-custom-plugin.test b/test-data/unit/check-custom-plugin.test index ee19113f000f..2ef52e5a3d33 100644 --- a/test-data/unit/check-custom-plugin.test +++ b/test-data/unit/check-custom-plugin.test @@ -991,3 +991,40 @@ class Cls: [file mypy.ini] \[mypy] plugins=/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=/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=/test-data/unit/plugins/function_decorator_hook.py diff --git a/test-data/unit/plugins/function_decorator_hook.py b/test-data/unit/plugins/function_decorator_hook.py new file mode 100644 index 000000000000..0efb30a7e89f --- /dev/null +++ b/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