[ty] Preserve short-circuit reachability for direct conditions

crates/ty_python_core/src/builder.rs

@@ -228,6 +228,40 @@ impl ConditionFlowSnapshot {
     }
 }
 
+/// Whether evaluation produces a result object or chooses a control-flow path.
+///
+/// In `Value` context, the enclosing code receives the expression's result object. For example,
+/// `result = x and y` produces `x` if `x` is falsy, or `y` otherwise. This also applies to expressions
+/// that return `bool`: the comparison in `result = x > 0` has value context.
+///
+/// In `Condition` context, the enclosing code only needs to know which branch to take. For example,
+/// CPython evaluates `if x and y` by testing `x` and, only if `x` is truthy, testing `y`. If `x` tests
+/// falsy, that one truthiness check is enough to skip the body: `x` is not tested again as the
+/// result of `x and y`.
+///
+/// This distinction matters when an operand's `__bool__` can change between calls:
+///
+/// ```python
+/// if x and False:      # A falsy x skips the body; a truthy x reaches False.
+///     ...              # Unreachable in either case.
+/// saved = x and False  # Can produce x after checking that it is falsy.
+/// if saved:            # Can call x.__bool__ again, which may now return True.
+///     ...              # Reachable.
+/// ```
+///
+/// The context propagates through `and`, `or`, `not`, and the branches of conditional expressions.
+/// Condition context does not propagate through calls or assignment expressions: in
+/// `if f(x and False)`, the call's result controls the branch, but its argument is evaluated in
+/// value context.
+#[derive(Clone, Copy, Debug, Default)]
+enum ExpressionContext {
+    /// Produce the expression's result object for the enclosing code to use.
+    #[default]
+    Value,
+    /// Choose the truthy or falsy control-flow path without preserving the result object.
+    Condition,
+}
+
 pub(super) struct SemanticIndexBuilder<'db, 'ast> {
     // Builder state
     db: &'db dyn Db,
@@ -253,6 +287,9 @@ pub(super) struct SemanticIndexBuilder<'db, 'ast> {
     has_future_annotations: bool,
     /// Whether we are currently visiting an `if TYPE_CHECKING` block.
     in_type_checking_block: bool,
+    /// How the next `visit_expr` consumes its expression. The visitor resets this to `Value`
+    /// before walking children; only short-circuit expressions propagate it to their operands.
+    next_expression_context: ExpressionContext,
 
     // Used for checking semantic syntax errors
     resolver_environment: ResolverEnvironment<'db>,
@@ -324,6 +361,7 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
 
             has_future_annotations: false,
             in_type_checking_block: false,
+            next_expression_context: ExpressionContext::Value,
 
             scopes: IndexVec::new(),
             place_tables: IndexVec::new(),
@@ -2153,12 +2191,16 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
         &mut self,
         predicate_node: &'ast ast::Expr,
     ) -> (PredicateOrLiteral<'db>, ScopedPredicateId) {
-        let predicate = self.build_predicate(predicate_node);
+        let predicate = self.build_predicate(predicate_node, ExpressionContext::Condition);
         let predicate_id = self.record_narrowing_constraint(predicate);
         (predicate, predicate_id)
     }
 
-    fn build_predicate(&mut self, predicate_node: &'ast ast::Expr) -> PredicateOrLiteral<'db> {
+    fn build_predicate(
+        &mut self,
+        predicate_node: &'ast ast::Expr,
+        context: ExpressionContext,
+    ) -> PredicateOrLiteral<'db> {
         // Some commonly used test expressions are eagerly evaluated as `true`
         // or `false` here for performance reasons. This list does not need to
         // be exhaustive. More complex expressions will still evaluate to the
@@ -2190,7 +2232,23 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
         match resolve_to_literal(predicate_node) {
             Some(literal) => PredicateOrLiteral::Literal(literal),
             None => PredicateOrLiteral::Predicate(Predicate {
-                node: PredicateNode::Expression(expression),
+                node: match (context, predicate_node) {
+                    (
+                        ExpressionContext::Condition,
+                        ast::Expr::BoolOp(_)
+                        | ast::Expr::If(_)
+                        | ast::Expr::UnaryOp(ast::ExprUnaryOp {
+                            op: ast::UnaryOp::Not,
+                            ..
+                        }),
+                    ) => PredicateNode::Condition(expression),
+                    (ExpressionContext::Condition, ast::Expr::Compare(compare))
+                        if compare.ops.len() > 1 =>
+                    {
+                        PredicateNode::ChainedComparisonCondition(expression)
+                    }
+                    _ => PredicateNode::Expression(expression),
+                },
                 is_positive: true,
             }),
         }
@@ -2249,7 +2307,9 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
                 let place_table = self.current_place_table();
 
                 match pred.node {
-                    PredicateNode::Expression(expression) => {
+                    PredicateNode::Expression(expression)
+                    | PredicateNode::Condition(expression)
+                    | PredicateNode::ChainedComparisonCondition(expression) => {
                         let expression_node = expression.node_ref(self.db).node(self.module);
                         let mut places = PossiblyNarrowedPlacesBuilder::new(self.db, place_table)
                             .expression(expression_node);
@@ -3008,7 +3068,7 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
     /// print(last)
     /// ```
     fn visit_comprehension_filter(&mut self, if_expr: &'ast ast::Expr) -> FlowSnapshot {
-        self.visit_expr(if_expr);
+        self.visit_expr_with_context(if_expr, ExpressionContext::Condition);
         let condition_flow_snapshot = self.flow_snapshot_for_condition(if_expr);
         let filtered_out = if let Some(snapshots) = condition_flow_snapshot.into_branches() {
             self.flow_restore(snapshots.truthy);
@@ -3273,13 +3333,18 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
             .get_or_init(|| source_text(self.db, self.file.file(self.db)))
     }
 
+    fn visit_expr_with_context(&mut self, expr: &'ast ast::Expr, context: ExpressionContext) {
+        self.next_expression_context = context;
+        self.visit_expr(expr);
+    }
+
     /// Visits a conditional expression without reserving its flow snapshots in every recursive
     /// expression-visitor frame. This matters for deeply nested expressions in unoptimized builds.
-    fn visit_if_expression(&mut self, node: &'ast ast::ExprIf) {
+    fn visit_if_expression(&mut self, node: &'ast ast::ExprIf, context: ExpressionContext) {
         let ast::ExprIf {
             body, test, orelse, ..
         } = node;
-        self.visit_expr(test);
+        self.visit_expr_with_context(test, ExpressionContext::Condition);
         let condition_flow_snapshot = self.flow_snapshot_for_condition(test);
         let falsy = if let Some(snapshots) = condition_flow_snapshot.into_branches() {
             self.flow_restore(snapshots.truthy);
@@ -3292,7 +3357,7 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
         let in_type_checking_block = self.in_type_checking_block;
         self.current_use_def_map_mut()
             .record_range_reachability(body.range(), in_type_checking_block);
-        self.visit_expr(body);
+        self.visit_expr_with_context(body, context);
         let post_body = self.flow_snapshot();
         self.flow_restore(falsy);
 
@@ -3301,12 +3366,12 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
         let in_type_checking_block = self.in_type_checking_block;
         self.current_use_def_map_mut()
             .record_range_reachability(orelse.range(), in_type_checking_block);
-        self.visit_expr(orelse);
+        self.visit_expr_with_context(orelse, context);
         self.flow_merge(post_body);
     }
 
     /// Keeps short-circuit flow snapshots out of the common recursive expression-visitor frame.
-    fn visit_bool_expression(&mut self, node: &'ast ast::ExprBoolOp) {
+    fn visit_bool_expression(&mut self, node: &'ast ast::ExprBoolOp, context: ExpressionContext) {
         let ast::ExprBoolOp { values, op, .. } = node;
         let mut snapshots = vec![];
         let mut reachability_constraints = vec![];
@@ -3321,7 +3386,7 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
             let in_type_checking_block = self.in_type_checking_block;
             self.current_use_def_map_mut()
                 .record_range_reachability(value.range(), in_type_checking_block);
-            self.visit_expr(value);
+            self.visit_expr_with_context(value, context);
 
             // Only non-final values can short-circuit this boolean operation. The final
             // value can still have its own outcome-specific flow if it is nested.
@@ -3330,7 +3395,7 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
                     value,
                 ));
                 let condition_flow_snapshots = self.take_condition_flow_snapshots(value);
-                let predicate = self.build_predicate(value);
+                let predicate = self.build_predicate(value, context);
                 let possibly_narrowed = self.compute_possibly_narrowed_places(&predicate);
                 let predicate_id = match op {
                     ast::BoolOp::And => self.add_predicate(predicate),
@@ -3870,9 +3935,9 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
                 // `msg` branch back into the following flow, since there is no way of getting out
                 // of that branch. Code after the assertion starts from the condition's truthy flow.
 
-                self.visit_expr(test);
+                self.visit_expr_with_context(test, ExpressionContext::Condition);
                 let condition_flow_snapshot = self.flow_snapshot_for_condition(test);
-                let predicate = self.build_predicate(test);
+                let predicate = self.build_predicate(test, ExpressionContext::Condition);
 
                 if msg.is_some()
                     || self
@@ -4043,7 +4108,7 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
                 }
             }
             ast::Stmt::If(node) => {
-                self.visit_expr(&node.test);
+                self.visit_expr_with_context(&node.test, ExpressionContext::Condition);
                 let condition_flow_snapshot = self.flow_snapshot_for_condition(&node.test);
                 let mut falsy = if let Some(snapshots) = condition_flow_snapshot.into_branches() {
                     self.flow_restore(snapshots.truthy);
@@ -4097,7 +4162,7 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
                     self.record_negated_reachability_constraint(last_reachability_constraint);
 
                     let next_falsy = if let Some(elif_test) = clause_test {
-                        self.visit_expr(elif_test);
+                        self.visit_expr_with_context(elif_test, ExpressionContext::Condition);
                         // A test expression is evaluated whether the branch is taken or not
                         let condition_flow_snapshot = self.flow_snapshot_for_condition(elif_test);
                         let next_falsy =
@@ -4180,7 +4245,7 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
 
                 // Visit the test expression after creating loop headers, so that loop-back values
                 // are visible.
-                self.visit_expr(test);
+                self.visit_expr_with_context(test, ExpressionContext::Condition);
                 let condition_flow_snapshot = self.flow_snapshot_for_condition(test);
 
                 // Take the pre_loop snapshot from the post-test fallback flow before restoring the
@@ -4558,7 +4623,7 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
                     // while the next case is reached through `!P || (P && !G)`. Save `P && !G`
                     // separately so it can be merged with the pattern-failure state after the body.
                     let match_success_guard_failure = case.guard.as_ref().map(|guard| {
-                        self.visit_expr(guard);
+                        self.visit_expr_with_context(guard, ExpressionContext::Condition);
                         let condition_flow_snapshot = self.flow_snapshot_for_condition(guard);
                         let falsy = if let Some(snapshots) = condition_flow_snapshot.into_branches()
                         {
@@ -5270,6 +5335,7 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> {
     }
 
     fn visit_expr(&mut self, expr: &'ast ast::Expr) {
+        let context = std::mem::take(&mut self.next_expression_context);
         self.with_semantic_checker(|semantic, context| semantic.visit_expr(expr, context));
 
         self.scopes_by_expression
@@ -5403,7 +5469,7 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> {
                 self.visit_expr(lambda.body.as_ref());
                 self.pop_scope();
             }
-            ast::Expr::If(node) => self.visit_if_expression(node),
+            ast::Expr::If(node) => self.visit_if_expression(node, context),
             ast::Expr::ListComp(
                 list_comprehension @ ast::ExprListComp {
                     elt, generators, ..
@@ -5470,7 +5536,14 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> {
                 self.record_exception_checkpoint();
             }
             ast::Expr::UnaryOp(unary) => {
-                walk_expr(self, expr);
+                self.visit_expr_with_context(
+                    &unary.operand,
+                    if unary.op == ast::UnaryOp::Not {
+                        context
+                    } else {
+                        ExpressionContext::Value
+                    },
+                );
                 self.record_exception_checkpoint_if(
                     unary.op != ast::UnaryOp::Not
                         || !Self::condition_evaluation_is_known_safe(&unary.operand),
@@ -5491,7 +5564,7 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> {
                     ));
                 }
             }
-            ast::Expr::BoolOp(node) => self.visit_bool_expression(node),
+            ast::Expr::BoolOp(node) => self.visit_bool_expression(node, context),
             ast::Expr::StringLiteral(_) => {
                 walk_expr(self, expr);
             }

crates/ty_python_core/src/predicate.rs

@@ -114,7 +114,16 @@ pub struct CallableAndCallExpr<'db> {
 
 #[derive(Clone, Copy, Debug, Hash, PartialEq, Eq, get_size2::GetSize, salsa::SalsaValue)]
 pub enum PredicateNode<'db> {
+    /// The truthiness of an expression's resulting value.
     Expression(Expression<'db>),
+    /// A boolean operation, `not`, or conditional expression evaluated directly as a condition.
+    ///
+    /// In `if x and False`, the truthy branch is unreachable. But after `y = x and False`,
+    /// `if y` may be truthy: it can call `x.__bool__` a second time and get a different result.
+    Condition(Expression<'db>),
+    /// A chained comparison evaluated directly as a condition. Its inferred truthiness is
+    /// available without walking the expression again.
+    ChainedComparisonCondition(Expression<'db>),
     /// Whether a context manager's exit return type allows an exception to be suppressed.
     ///
     /// Resolved during type inference because the context manager's type is unavailable during

crates/ty_python_semantic/resources/mdtest/annotations/annotated.md

@@ -37,6 +37,15 @@ def convert(value: "Annotated[str, dict(**{'name': 'value'})]") -> "Annotated[in
     return 1
 ```
 
+Conditional expressions are also valid metadata and do not affect the annotated type.
+
+```py
+def flag() -> bool:
+    return True
+
+conditional_value: "Annotated[int, 1 if flag() else 2]" = 1
+```
+
 ## Inside `type[...]`
 
 `Annotated` can wrap a class or specialized generic class inside `type[...]` without changing the