type check function decorators (#428)

docs/changelog.md

@@ -2,6 +2,7 @@
 
 ## Unreleased
 
+- Type check function decorators (#428)
 - Handle `NoReturn` in `async def` functions (#427)
 - Support PEP 673 (`typing_extensions.Self`) (#423)
 - Updates for compatibility with recent changes in typeshed (#421):

pyanalyze/functions.py

@@ -66,10 +66,10 @@ class FunctionInfo:
     is_overload: bool  # typing.overload or pyanalyze.extensions.overload
     is_evaluated: bool  # @pyanalyze.extensions.evaluated
     is_abstractmethod: bool  # has @abstractmethod
-    # a list of pairs of (decorator function, applied decorator function). These are different
-    # for decorators that take arguments, like @asynq(): the first element will be the asynq
-    # function and the second will be the result of calling asynq().
-    decorators: List[Tuple[Value, Value]]
+    # a list of tuples of (decorator function, applied decorator function, AST node). These are
+    # different for decorators that take arguments, like @asynq(): the first element will be the
+    # asynq function and the second will be the result of calling asynq().
+    decorators: List[Tuple[Value, Value, ast.AST]]
     node: FunctionNode
     params: Sequence[ParamInfo]
     return_annotation: Optional[Value]
@@ -162,7 +162,7 @@ def compute_function_info(
                     # @async_proxy(pure=True) is a noop, so don't treat it specially
                     if not any(kw.arg == "pure" for kw in decorator.keywords):
                         async_kind = AsyncFunctionKind.async_proxy
-            decorators.append((callee, decorator_value))
+            decorators.append((callee, decorator_value, decorator))
         else:
             decorator_value = ctx.visit_expression(decorator)
             if decorator_value == KnownValue(classmethod):
@@ -179,7 +179,7 @@ def compute_function_info(
                 is_abstractmethod = True
             elif decorator_value == KnownValue(evaluated):
                 is_evaluated = True
-            decorators.append((decorator_value, decorator_value))
+            decorators.append((decorator_value, decorator_value, decorator))
     params = compute_parameters(
         node,
         enclosing_class,
@@ -323,7 +323,7 @@ def compute_value_of_function(
         has_return_annotation=info.return_annotation is not None,
     )
     val = CallableValue(sig, types.FunctionType)
-    for unapplied, decorator in reversed(info.decorators):
+    for unapplied, decorator, node in reversed(info.decorators):
         # Special case asynq.asynq until we can create the type automatically
         if unapplied == KnownValue(asynq.asynq) and isinstance(val, CallableValue):
             sig = replace(val.signature, is_asynq=True)
@@ -332,7 +332,5 @@ def compute_value_of_function(
         allow_call = isinstance(
             unapplied, KnownValue
         ) and SafeDecoratorsForNestedFunctions.contains(unapplied.val, ctx.options)
-        val = ctx.check_call(
-            info.node, decorator, [Composite(val)], allow_call=allow_call
-        )
+        val = ctx.check_call(node, decorator, [Composite(val)], allow_call=allow_call)
     return val

pyanalyze/name_check_visitor.py

@@ -1578,8 +1578,9 @@ def visit_FunctionDef(self, node: FunctionDefNode) -> Value:
                 node, error_code=ErrorCode.missing_return_annotation
             )
 
+        computed_function = compute_value_of_function(info, self)
         if potential_function is None:
-            val = compute_value_of_function(info, self)
+            val = computed_function
         else:
             val = KnownValue(potential_function)
         if not info.is_overload and not info.is_evaluated:

pyanalyze/stacked_scopes.py

@@ -1109,7 +1109,7 @@ def get_origin(
         return self._resolve_origin(definers)
 
     @contextlib.contextmanager
-    def subscope(self) -> Iterable[SubScope]:
+    def subscope(self) -> Iterator[SubScope]:
         """Create a new subscope, to be used for conditional branches."""
         # Ignore LEAVES_SCOPE if it's already there, so that we type check code after the
         # assert False correctly. Without this, test_after_assert_false fails.
@@ -1127,7 +1127,7 @@ def subscope(self) -> Iterable[SubScope]:
             yield new_name_to_nodes
 
     @contextlib.contextmanager
-    def loop_scope(self) -> Iterable[SubScope]:
+    def loop_scope(self) -> Iterator[SubScope]:
         loop_scopes = []
         with self.subscope() as main_scope:
             loop_scopes.append(main_scope)
@@ -1276,7 +1276,7 @@ def add_scope(
         scope_type: ScopeType,
         scope_node: Node,
         scope_object: Optional[object] = None,
-    ) -> Iterable[None]:
+    ) -> Iterator[None]:
         """Context manager that adds a scope of this type to the top of the stack."""
         if scope_type is ScopeType.function_scope:
             scope = FunctionScope(

pyanalyze/test_functions.py

@@ -164,3 +164,14 @@ def capybara():
 
             fun4 = lambda a, b, c: a if c else b
             assert_is_value(fun4(1, 2, 3), KnownValue(1) | KnownValue(2))
+
+
+class TestDecorators(TestNameCheckVisitorBase):
+    @assert_passes()
+    def test_applied(self) -> None:
+        def bad_deco(x: int) -> str:
+            return "x"
+
+        @bad_deco  # E: incompatible_argument
+        def capybara():
+            pass