Fix error code on "Maybe you forgot to use await" note (#16203)

Fixes #16202

Author: JelleZijlstra

mypy/checker.py

@@ -6237,7 +6237,7 @@ def check_subtype(
                 assert call is not None
                 if not is_subtype(subtype, call, options=self.options):
                     self.msg.note_call(supertype, call, context, code=msg.code)
-        self.check_possible_missing_await(subtype, supertype, context)
+        self.check_possible_missing_await(subtype, supertype, context, code=msg.code)
         return False
 
     def get_precise_awaitable_type(self, typ: Type, local_errors: ErrorWatcher) -> Type | None:
@@ -6271,7 +6271,7 @@ def checking_await_set(self) -> Iterator[None]:
             self.checking_missing_await = False
 
     def check_possible_missing_await(
-        self, subtype: Type, supertype: Type, context: Context
+        self, subtype: Type, supertype: Type, context: Context, code: ErrorCode | None
     ) -> None:
         """Check if the given type becomes a subtype when awaited."""
         if self.checking_missing_await:
@@ -6285,7 +6285,7 @@ def check_possible_missing_await(
                 aw_type, supertype, context, msg=message_registry.INCOMPATIBLE_TYPES
             ):
                 return
-        self.msg.possible_missing_await(context)
+        self.msg.possible_missing_await(context, code)
 
     def contains_none(self, t: Type) -> bool:
         t = get_proper_type(t)

mypy/checkexpr.py

@@ -2563,7 +2563,7 @@ def check_arg(
                 original_caller_type, callee_type, context, code=code
             )
             if not self.msg.prefer_simple_messages():
-                self.chk.check_possible_missing_await(caller_type, callee_type, context)
+                self.chk.check_possible_missing_await(caller_type, callee_type, context, code)
 
     def check_overload_call(
         self,

mypy/checkmember.py

@@ -272,11 +272,11 @@ def report_missing_attribute(
     mx: MemberContext,
     override_info: TypeInfo | None = None,
 ) -> Type:
-    res_type = mx.msg.has_no_attr(original_type, typ, name, mx.context, mx.module_symbol_table)
+    error_code = mx.msg.has_no_attr(original_type, typ, name, mx.context, mx.module_symbol_table)
     if not mx.msg.prefer_simple_messages():
         if may_be_awaitable_attribute(name, typ, mx, override_info):
-            mx.msg.possible_missing_await(mx.context)
-    return res_type
+            mx.msg.possible_missing_await(mx.context, error_code)
+    return AnyType(TypeOfAny.from_error)
 
 
 # The several functions that follow implement analyze_member_access for various

mypy/messages.py

@@ -355,7 +355,7 @@ def has_no_attr(
         member: str,
         context: Context,
         module_symbol_table: SymbolTable | None = None,
-    ) -> Type:
+    ) -> ErrorCode | None:
         """Report a missing or non-accessible member.
 
         original_type is the top-level type on which the error occurred.
@@ -370,44 +370,49 @@ def has_no_attr(
         directly available on original_type
 
         If member corresponds to an operator, use the corresponding operator
-        name in the messages. Return type Any.
+        name in the messages. Return the error code that was produced, if any.
         """
         original_type = get_proper_type(original_type)
         typ = get_proper_type(typ)
 
         if isinstance(original_type, Instance) and original_type.type.has_readable_member(member):
             self.fail(f'Member "{member}" is not assignable', context)
+            return None
         elif member == "__contains__":
             self.fail(
                 f"Unsupported right operand type for in ({format_type(original_type, self.options)})",
                 context,
                 code=codes.OPERATOR,
             )
+            return codes.OPERATOR
         elif member in op_methods.values():
             # Access to a binary operator member (e.g. _add). This case does
             # not handle indexing operations.
             for op, method in op_methods.items():
                 if method == member:
                     self.unsupported_left_operand(op, original_type, context)
-                    break
+                    return codes.OPERATOR
         elif member == "__neg__":
             self.fail(
                 f"Unsupported operand type for unary - ({format_type(original_type, self.options)})",
                 context,
                 code=codes.OPERATOR,
             )
+            return codes.OPERATOR
         elif member == "__pos__":
             self.fail(
                 f"Unsupported operand type for unary + ({format_type(original_type, self.options)})",
                 context,
                 code=codes.OPERATOR,
             )
+            return codes.OPERATOR
         elif member == "__invert__":
             self.fail(
                 f"Unsupported operand type for ~ ({format_type(original_type, self.options)})",
                 context,
                 code=codes.OPERATOR,
             )
+            return codes.OPERATOR
         elif member == "__getitem__":
             # Indexed get.
             # TODO: Fix this consistently in format_type
@@ -418,12 +423,14 @@ def has_no_attr(
                     ),
                     context,
                 )
+                return None
             else:
                 self.fail(
                     f"Value of type {format_type(original_type, self.options)} is not indexable",
                     context,
                     code=codes.INDEX,
                 )
+                return codes.INDEX
         elif member == "__setitem__":
             # Indexed set.
             self.fail(
@@ -433,19 +440,22 @@ def has_no_attr(
                 context,
                 code=codes.INDEX,
             )
+            return codes.INDEX
         elif member == "__call__":
             if isinstance(original_type, Instance) and (
                 original_type.type.fullname == "builtins.function"
             ):
                 # "'function' not callable" is a confusing error message.
                 # Explain that the problem is that the type of the function is not known.
                 self.fail("Cannot call function of unknown type", context, code=codes.OPERATOR)
+                return codes.OPERATOR
             else:
                 self.fail(
                     message_registry.NOT_CALLABLE.format(format_type(original_type, self.options)),
                     context,
                     code=codes.OPERATOR,
                 )
+                return codes.OPERATOR
         else:
             # The non-special case: a missing ordinary attribute.
             extra = ""
@@ -501,6 +511,7 @@ def has_no_attr(
                         context,
                         code=codes.ATTR_DEFINED,
                     )
+                return codes.ATTR_DEFINED
             elif isinstance(original_type, UnionType):
                 # The checker passes "object" in lieu of "None" for attribute
                 # checks, so we manually convert it back.
@@ -518,6 +529,7 @@ def has_no_attr(
                     context,
                     code=codes.UNION_ATTR,
                 )
+                return codes.UNION_ATTR
             elif isinstance(original_type, TypeVarType):
                 bound = get_proper_type(original_type.upper_bound)
                 if isinstance(bound, UnionType):
@@ -531,6 +543,7 @@ def has_no_attr(
                         context,
                         code=codes.UNION_ATTR,
                     )
+                    return codes.UNION_ATTR
             else:
                 self.fail(
                     '{} has no attribute "{}"{}'.format(
@@ -539,7 +552,8 @@ def has_no_attr(
                     context,
                     code=codes.ATTR_DEFINED,
                 )
-        return AnyType(TypeOfAny.from_error)
+                return codes.ATTR_DEFINED
+        return None
 
     def unsupported_operand_types(
         self,
@@ -1107,8 +1121,8 @@ def unpacking_strings_disallowed(self, context: Context) -> None:
     def type_not_iterable(self, type: Type, context: Context) -> None:
         self.fail(f"{format_type(type, self.options)} object is not iterable", context)
 
-    def possible_missing_await(self, context: Context) -> None:
-        self.note('Maybe you forgot to use "await"?', context)
+    def possible_missing_await(self, context: Context, code: ErrorCode | None) -> None:
+        self.note('Maybe you forgot to use "await"?', context, code=code)
 
     def incompatible_operator_assignment(self, op: str, context: Context) -> None:
         self.fail(f"Result type of {op} incompatible in assignment", context)

test-data/unit/check-async-await.test

@@ -165,6 +165,33 @@ async def f() -> None:
 [out]
 main:4: error: "List[int]" has no attribute "__aiter__" (not async iterable)
 
+[case testAsyncForErrorNote]
+
+from typing import AsyncIterator, AsyncGenerator
+async def g() -> AsyncGenerator[str, None]:
+    pass
+
+async def f() -> None:
+    async for x in g():
+        pass
+[builtins fixtures/async_await.pyi]
+[typing fixtures/typing-async.pyi]
+[out]
+main:7: error: "Coroutine[Any, Any, AsyncGenerator[str, None]]" has no attribute "__aiter__" (not async iterable)
+main:7: note: Maybe you forgot to use "await"?
+
+[case testAsyncForErrorCanBeIgnored]
+
+from typing import AsyncIterator, AsyncGenerator
+async def g() -> AsyncGenerator[str, None]:
+    pass
+
+async def f() -> None:
+    async for x in g():  # type: ignore[attr-defined]
+        pass
+[builtins fixtures/async_await.pyi]
+[typing fixtures/typing-async.pyi]
+
 [case testAsyncForTypeComments]
 
 from typing import AsyncIterator, Union