[ty] Add await fixes for redundant conditions (#28169)

`redundant-condition` already catches conditions that test a coroutine
object instead of its result. Add an unsafe fix that awaits the
expression, inserting parentheses when needed to preserve precedence.
For an uncalled function or method known to return a coroutine, extend
the existing call fix to call and await it.

Only suggest `await` where it is syntactically valid, accounting for
annotations, nested scopes, comprehensions, Python-version differences,
and top-level await in notebooks. A return type of `Any`, `Unknown`, or
`Never` does not by itself justify adding `await`.

Stacked on #28168.

crates/ty_python_core/src/builder.rs

@@ -300,6 +300,8 @@ pub(super) struct SemanticIndexBuilder<'db, 'ast> {
     generator_functions: FxHashSet<FileScopeId>,
     /// Hashset of all [`FileScopeId`]s that correspond to asynchronous comprehensions.
     async_comprehensions: FxHashSet<FileScopeId>,
+    /// Node indices of syntactic annotation roots.
+    annotations: Vec<NodeIndex>,
     /// Snapshots of enclosing-scope place states visible from nested scopes.
     enclosing_snapshots: FxHashMap<EnclosingSnapshotKey, ScopedEnclosingSnapshotId>,
     /// Errors collected by the `semantic_checker`.
@@ -356,6 +358,7 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
             imported_modules: FxHashSet::default(),
             generator_functions: FxHashSet::default(),
             async_comprehensions: FxHashSet::default(),
+            annotations: Vec::new(),
 
             enclosing_snapshots: FxHashMap::default(),
 
@@ -3343,6 +3346,8 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
 
         let mut semantic_syntax_errors = self.semantic_syntax_errors.into_inner();
         semantic_syntax_errors.shrink_to_fit();
+        // Node indices follow source order, while semantic visitation may not.
+        self.annotations.sort_unstable();
         let uses_by_collection = FrozenMap::from_entries(
             self.uses_by_collection
                 .into_iter()
@@ -3381,6 +3386,7 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
             semantic_syntax_errors,
             generator_functions: FrozenSet::from(self.generator_functions),
             async_comprehensions: FrozenSet::from(self.async_comprehensions),
+            annotations: self.annotations.into_boxed_slice(),
             narrowing_alias_predicates: FrozenMap::from(self.alias_predicates),
         }
     }
@@ -4313,7 +4319,7 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
                 // not discard the declared type. The value is still bound only after the RHS
                 // completes, so a handler can observe an earlier binding (or an unbound name).
                 let pending = self.begin_annotated_assignment(node);
-                self.visit_expr(&node.annotation);
+                self.visit_annotation(&node.annotation);
                 if let Some(value) = &node.value {
                     self.visit_expr(value);
                     if self.is_method_or_eagerly_executed_in_method().is_some() {
@@ -5503,6 +5509,11 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
 }
 
 impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> {
+    fn visit_annotation(&mut self, expression: &'ast ast::Expr) {
+        self.annotations.push(expression.node_index().load());
+        self.visit_expr(expression);
+    }
+
     fn visit_stmt(&mut self, stmt: &'ast ast::Stmt) {
         self.push_statement(CurrentStatement::default());
         self.visit_stmt_impl(stmt);

crates/ty_python_core/src/lib.rs

@@ -6,12 +6,12 @@ use ruff_python_ast as ast;
 use std::iter::{FusedIterator, once};
 use std::sync::Arc;
 
-use ruff_db::parsed::parsed_module;
+use ruff_db::parsed::{ParsedModuleRef, parsed_module};
 
 use ruff_index::{FrozenIndexVec, IndexSlice};
 use ruff_python_ast::{HasNodeIndex, NodeIndex};
 use ruff_python_parser::semantic_errors::SemanticSyntaxError;
-use ruff_text_size::TextRange;
+use ruff_text_size::{Ranged, TextRange};
 use rustc_hash::{FxHashMap, FxHashSet};
 use salsa::plumbing::AsId;
 use smallvec::SmallVec;
@@ -345,6 +345,11 @@ pub struct SemanticIndex<'db> {
     /// Set of all asynchronous comprehensions in this file.
     async_comprehensions: FrozenSet<FileScopeId>,
 
+    /// Node indices of syntactic annotation roots, sorted in source order. Annotations cannot
+    /// contain other syntactic annotations, so their ranges do not overlap. Storing only indices
+    /// avoids duplicating ranges and scopes already available from the AST and expression scope map.
+    annotations: Box<[NodeIndex]>,
+
     /// Narrowing alias metadata for predicate leaf names.
     /// When a predicate references an alias variable (e.g., `is_none` from `is_none = x is None`),
     /// the alias Name node is mapped to its aliased expression for constraint-generation time.
@@ -416,6 +421,40 @@ impl<'db> SemanticIndex<'db> {
         self.scopes_by_expression.try_get(expression)
     }
 
+    /// Returns the scope enclosing the syntactic annotation containing `expression`, if any.
+    ///
+    /// `module` must correspond to the same file and revision as this index.
+    ///
+    /// ```python
+    /// from typing import Annotated
+    ///
+    /// def example(items: list[int]) -> None:
+    ///     result: Annotated[int, (item for item in items)]
+    /// ```
+    ///
+    /// The yielded `item` belongs to the generator's scope, and the iterable `items` belongs to
+    /// `example`'s scope. This method returns `example`'s scope for both, because both expressions
+    /// occur in the annotation on `result`.
+    pub fn annotation_parent_scope_id(
+        &self,
+        module: &ParsedModuleRef,
+        expression: &impl Ranged,
+    ) -> Option<FileScopeId> {
+        let index = self
+            .annotations
+            .partition_point(|index| module.get_by_index(*index).start() <= expression.start())
+            .checked_sub(1)?;
+        let ast::AnyRootNodeRef::Expr(annotation) = module.get_by_index(self.annotations[index])
+        else {
+            return None;
+        };
+        if annotation.range().contains_range(expression.range()) {
+            self.try_expression_scope_id(annotation)
+        } else {
+            None
+        }
+    }
+
     /// Returns the [`Scope`] of the `expression`'s enclosing scope.
     #[allow(unused)]
     #[track_caller]
@@ -1916,23 +1955,38 @@ class C[T]:
 
     #[test]
     fn expression_scope() {
-        let TestCase { db, file } = test_case("x = 1;\ndef test():\n  y = 4");
+        // Annotation tracking preserves lexical scope lookup, including for attribute identifiers
+        // inside annotations, which are not registered separately in the expression scope map.
+        // `annotation_parent_scope_id` returns the module scope for `module.Type`, the annotation
+        // on `x`. It returns `None` for `x` itself and for `y`, since neither is inside an annotation.
+        let TestCase { db, file } = test_case("x: module.Type = 1;\ndef test():\n  y = 4");
 
         let index = semantic_index(&db, program_file(&db, file));
         let module = parsed_module(&db, program_file(&db, file).python_file(&db)).load(&db);
         let ast = module.syntax();
 
-        let x_stmt = ast.body[0].as_assign_stmt().unwrap();
-        let x = &x_stmt.targets[0];
+        let x_stmt = ast.body[0].as_ann_assign_stmt().unwrap();
+        let x = x_stmt.target.as_ref();
 
         assert_eq!(index.expression_scope(x).kind(), ScopeKind::Module);
         assert_eq!(index.expression_scope_id(x), FileScopeId::global());
+        assert_matches!(
+            x_stmt.annotation.as_ref(),
+            ast::Expr::Attribute(attribute)
+                if index.try_expression_scope_id(&attribute.attr) == Some(FileScopeId::global())
+        );
+        assert_eq!(index.annotation_parent_scope_id(&module, x), None);
+        assert_eq!(
+            index.annotation_parent_scope_id(&module, x_stmt.annotation.as_ref()),
+            Some(FileScopeId::global())
+        );
 
         let def = ast.body[1].as_function_def_stmt().unwrap();
         let y_stmt = def.body[0].as_assign_stmt().unwrap();
         let y = &y_stmt.targets[0];
 
         assert_eq!(index.expression_scope(y).kind(), ScopeKind::Function);
+        assert_eq!(index.annotation_parent_scope_id(&module, y), None);
     }
 
     #[test]

crates/ty_python_semantic/src/types/infer/builder/redundant_conditions/diagnostic.rs

@@ -8,11 +8,15 @@ use ruff_db::{
     source::source_text,
 };
 use ruff_diagnostics::{Applicability, Edit, Fix};
-use ruff_python_ast as ast;
+use ruff_python_ast::{self as ast, PythonVersion};
 use ruff_source_file::LineRanges;
 use ruff_text_size::Ranged;
 use ty_module_resolver::{SearchPath, file_to_module};
-use ty_python_core::{Truthiness, definition::DefinitionKind};
+use ty_python_core::{
+    Truthiness,
+    definition::DefinitionKind,
+    scope::{NodeWithScopeKind, ScopeKind},
+};
 
 use crate::{
     SemanticModel,
@@ -216,16 +220,39 @@ impl<'db> TypeInferenceBuilder<'db, '_> {
                 let mut diagnostic =
                     builder.into_diagnostic(format_args!("{function} is always truthy"));
 
-                // Add a suggestion and fix that they might have meant to call this function.
+                // Add a suggestion and fix that they might have meant to call (and possibly
+                // also await) this function.
                 //
                 // It's true that calling the function might not actually fix this diagnostic
                 // if the function returns something that is always truthy. They still probably
                 // meant to call the function, though, so it's still a useful suggestion/fix!
 
-                diagnostic.set_primary_annotation_message(format_args!(
-                    "Did you mean to call this {}?",
-                    function.kind()
-                ));
+                // A coroutine return type establishes that calling and awaiting the function
+                // is appropriate. `Any`, `Unknown`, and `Never` do not establish this, even
+                // though they are assignable to `CoroutineType`.
+                // Use the top materialization so the unspecified generic arguments do not
+                // prevent concrete coroutine types from being subtypes.
+                let coroutine = KnownClass::CoroutineType
+                    .to_instance(db, env)
+                    .top_materialization(db, env);
+
+                let is_awaitable_coro_function = self.can_await_here(test)
+                    && function.signature().iter().all(|signature| {
+                        !signature.return_ty.is_equivalent_to(db, env, Type::Never)
+                            && signature.return_ty.is_subtype_of(db, env, coroutine)
+                    });
+
+                let kind = function.kind();
+
+                if is_awaitable_coro_function {
+                    diagnostic.set_primary_annotation_message(format_args!(
+                        "Did you mean to `await` and call this {kind}?",
+                    ));
+                } else {
+                    diagnostic.set_primary_annotation_message(format_args!(
+                        "Did you mean to call this {kind}?"
+                    ));
+                }
 
                 if matches!(test, ast::Expr::Name(_) | ast::Expr::Attribute(_)) {
                     let (call, applicability) = if function.signature().has_parameters() {
@@ -235,7 +262,16 @@ impl<'db> TypeInferenceBuilder<'db, '_> {
                     };
                     let call_edit = Edit::insertion(call.to_string(), test.end());
 
-                    diagnostic.set_fix(Fix::applicable_edit(call_edit, applicability));
+                    let fix = if is_awaitable_coro_function {
+                        Fix::applicable_edits(
+                            Edit::insertion("await ".to_string(), test.start()),
+                            [call_edit],
+                            applicability,
+                        )
+                    } else {
+                        Fix::applicable_edit(call_edit, applicability)
+                    };
+                    diagnostic.set_fix(fix);
                 }
 
                 diagnostic
@@ -399,7 +435,20 @@ impl<'db> TypeInferenceBuilder<'db, '_> {
                 let mut diagnostic = builder.into_diagnostic("Condition is always truthy");
                 add_always_truthy_concise_message(&mut diagnostic);
                 annotate_inferred_type(&mut diagnostic);
-                if let Type::NominalInstance(instance) = test_type {
+                if test_type.try_await(db, env).is_ok() && self.can_await_here(test) {
+                    diagnostic.help("Did you mean to `await` this expression?");
+
+                    let fix = if test.precedence() <= ast::OperatorPrecedence::Await {
+                        Fix::unsafe_edits(
+                            Edit::insertion("await (".to_string(), test.start()),
+                            [Edit::insertion(")".to_string(), test.end())],
+                        )
+                    } else {
+                        Fix::unsafe_edit(Edit::insertion("await ".to_string(), test.start()))
+                    };
+
+                    diagnostic.set_fix(fix);
+                } else if let Type::NominalInstance(instance) = test_type {
                     let class = instance.class(db, env);
                     if class.is_final(db)
                         && !class.is_known(db, KnownClass::CoroutineType)
@@ -729,4 +778,70 @@ impl<'db> TypeInferenceBuilder<'db, '_> {
             }
         }
     }
+
+    /// Returns `true` if adding `await` at `expression` would produce valid Python.
+    ///
+    /// Accounts for asynchronous functions, notebook cells, annotation restrictions, enclosing
+    /// scopes, and the different scoping behavior of comprehensions and generator expressions.
+    fn can_await_here(&self, expression: &ast::Expr) -> bool {
+        let Some(expression_scope) = self.index.try_expression_scope_id(expression) else {
+            return false;
+        };
+        let annotation_parent_scope = self
+            .index
+            .annotation_parent_scope_id(self.module(), expression);
+
+        let db = self.db();
+
+        let mut in_eager_comprehension = false;
+
+        for (scope_id, scope) in self.index.ancestor_scopes(expression_scope) {
+            // The first iterable of a comprehension stays in the annotation's enclosing scope.
+            // Eager comprehensions also inherit the restriction, but a generator body can allow
+            // `await` before we reach the scope enclosing its annotation.
+            // Conservatively reject annotations on every Python version, even though some allow
+            // `await` before Python 3.14 without `from __future__ import annotations`. Avoiding
+            // invalid syntax matters more than offering every possible fix in this rare context.
+            if Some(scope_id) == annotation_parent_scope {
+                return false;
+            }
+
+            // Before Python 3.11, awaiting in a nested list, set, or dict comprehension cannot
+            // implicitly make its containing comprehension or generator expression asynchronous.
+            if in_eager_comprehension
+                && scope.kind() == ScopeKind::Comprehension
+                && self.program_environment().python_version(db) < PythonVersion::PY311
+                && !scope_id.is_async_comprehension(self.index)
+            {
+                return false;
+            }
+
+            match scope.node() {
+                NodeWithScopeKind::Function(function) => {
+                    return function.node(self.module()).is_async;
+                }
+                NodeWithScopeKind::Lambda(_)
+                | NodeWithScopeKind::Class(_)
+                | NodeWithScopeKind::ClassTypeParameters(_)
+                | NodeWithScopeKind::FunctionTypeParameters(_)
+                | NodeWithScopeKind::TypeAliasTypeParameters(_)
+                | NodeWithScopeKind::TypeAlias(_) => {
+                    return false;
+                }
+                NodeWithScopeKind::GeneratorExpression(_) => {
+                    return true;
+                }
+                NodeWithScopeKind::Module => {
+                    return source_text(db, self.file()).is_notebook();
+                }
+                NodeWithScopeKind::DictComprehension(_)
+                | NodeWithScopeKind::ListComprehension(_)
+                | NodeWithScopeKind::SetComprehension(_) => {
+                    in_eager_comprehension = true;
+                }
+            }
+        }
+
+        false
+    }
 }