Improve annotations of calls and assignments

Author: JukkaL

mypyc/annotate.py

@@ -22,14 +22,20 @@
     ForStmt,
     FuncDef,
     GeneratorExpr,
+    IndexExpr,
     LambdaExpr,
     MemberExpr,
     MypyFile,
+    NamedTupleExpr,
     NameExpr,
+    NewTypeExpr,
     Node,
+    OpExpr,
     RefExpr,
     TupleExpr,
+    TypedDictExpr,
     TypeInfo,
+    TypeVarExpr,
     Var,
     WithStmt,
 )
@@ -204,7 +210,7 @@ def function_annotations(func_ir: FuncIR, tree: MypyFile) -> dict[int, list[Anno
                 elif name == "PyObject_VectorcallMethod":
                     method_name = get_str_literal(op.args[0])
                     if method_name:
-                        ann = f'Call non-native method "{method_name}".'
+                        ann = f'Call non-native method "{method_name}" (it may be defined in a non-native class, or decorated).'
                     else:
                         ann = "Dynamic method call."
                 elif name in op_hints:
@@ -300,6 +306,10 @@ def visit_class_def(self, o: ClassDef, /) -> None:
                 if isinstance(s, AssignmentStmt):
                     # Don't complain about attribute initializers
                     self.ignored_lines.add(s.line)
+                elif isinstance(s, Decorator):
+                    # Don't complain about decorator definitions that generate some
+                    # dynamic operations. This is a bit heavy-handed.
+                    self.ignored_lines.add(s.func.line)
 
     def visit_with_stmt(self, o: WithStmt, /) -> None:
         for expr in o.expr:
@@ -317,10 +327,25 @@ def visit_with_stmt(self, o: WithStmt, /) -> None:
                             f'"{node.name}" uses @contextmanager, which is slow '
                             + "in compiled code. Use a native class with "
                             + '"__enter__" and "__exit__" methods instead.',
-                            priority=2,
+                            priority=3,
                         )
         super().visit_with_stmt(o)
 
+    def visit_assignment_stmt(self, o: AssignmentStmt, /) -> None:
+        special_form = False
+        if self.func_depth == 0:
+            analyzed = o.rvalue
+            if isinstance(o.rvalue, (CallExpr, IndexExpr, OpExpr)):
+                analyzed = o.rvalue.analyzed
+            if o.is_alias_def or isinstance(
+                analyzed, (TypeVarExpr, NamedTupleExpr, TypedDictExpr, NewTypeExpr)
+            ):
+                special_form = True
+            if special_form:
+                # TODO: Ignore all lines if multi-line
+                self.ignored_lines.add(o.line)
+        super().visit_assignment_stmt(o)
+
     def visit_name_expr(self, o: NameExpr, /) -> None:
         if ann := stdlib_hints.get(o.fullname):
             self.annotate(o, ann)
@@ -355,8 +380,14 @@ def visit_call_expr(self, o: CallExpr, /) -> None:
                     + "constructing an instance is slow.",
                     2,
                 )
-
-            print(o.callee.node.fullname, info in self.mapper.type_to_ir)
+        elif isinstance(o.callee, RefExpr) and isinstance(o.callee.node, Decorator):
+            decorator = o.callee.node
+            if self.mapper.is_native_ref_expr(o.callee):
+                self.annotate(
+                    o,
+                    f'Calling a decorated function ("{decorator.name}") is inefficient, even if it\'s native.',
+                    2,
+                )
 
     def check_isinstance_arg(self, arg: Expression) -> None:
         if isinstance(arg, RefExpr):

mypyc/test-data/annotate-basic.test

@@ -109,11 +109,11 @@ def f(x):
 from typing import Any
 
 def f1(x):
-    return x.foo()  # A: Call non-native method "foo".
+    return x.foo()  # A: Call non-native method "foo" (it may be defined in a non-native class, or decorated).
 
 def f2(x: Any) -> None:
-    x.foo(1)  # A: Call non-native method "foo".
-    x.foo(a=1)  # A: Call non-native method "foo".
+    x.foo(1)  # A: Call non-native method "foo" (it may be defined in a non-native class, or decorated).
+    x.foo(a=1)  # A: Call non-native method "foo" (it may be defined in a non-native class, or decorated).
     t = (1, 'x')
     x.foo(*t)  # A: Get non-native attribute "foo". Generic call operation.
     d = {"a": 1}
@@ -186,7 +186,7 @@ def startswith(x: str) -> bool:
     return x.startswith('foo')
 
 def islower(x: str) -> bool:
-    return x.islower()  # A: Call non-native method "islower".
+    return x.islower()  # A: Call non-native method "islower" (it may be defined in a non-native class, or decorated).
 
 [case testAnnotateSpecificStdlibFeatures]
 import functools
@@ -351,7 +351,7 @@ from nonnative import C
 
 def f1() -> None:
     c = C()  # A: Creating an instance of non-native class "C" is slow.
-    c.foo()  # A: Call non-native method "foo".
+    c.foo()  # A: Call non-native method "foo" (it may be defined in a non-native class, or decorated).
 
 class NT(NamedTuple):
     x: int
@@ -396,3 +396,40 @@ def f1(x, s: str):
 
 def f2(x, s: str):
     setattr(x, s, None)  # A: Dynamic attribute set.
+
+[case testAnnotateSpecialAssignments]
+from typing import TypeVar, NamedTuple, List, TypedDict, NewType
+
+# Even though these are slow, we don't complain about them since there is generally
+# no better way (and at module top level these are very unlikely to be bottlenecks)
+A = List[int]
+T = TypeVar("T", bound=List[int])
+NT = NamedTuple("NT", [("x", List[int])])
+TD = TypedDict("TD", {"x": List[int]})
+New = NewType("New", List[int])
+[typing fixtures/typing-full.pyi]
+
+[case testAnnotateCallDecoratedNativeFunctionOrMethod]
+from typing import TypeVar, Callable, Any
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+def mydeco(f: F) -> F:
+    return f
+
+@mydeco
+def d(x: int) -> int:
+    return x
+
+def f1() -> int:
+    return d(1)  # A: Calling a decorated function ("d") is inefficient, even if it's native.
+
+class C:
+    @mydeco
+    def d(self) -> None:
+        pass
+
+
+def f2() -> None:
+    c = C()
+    c.d()  # A: Call non-native method "d" (it may be defined in a non-native class, or decorated).

mypyc/test/test_annotate.py

@@ -50,7 +50,9 @@ def run_case(self, testcase: DataDrivenTestCase) -> None:
             else:
                 annotations = generate_annotations("native.py", tree, ir, type_map, mapper)
                 actual = []
-                for line_num, line_anns in annotations.annotations.items():
+                for line_num, line_anns in sorted(
+                    annotations.annotations.items(), key=lambda it: it[0]
+                ):
                     anns = get_max_prio(line_anns)
                     str_anns = [a.message for a in anns]
                     s = " ".join(str_anns)