[ty] Correctly handle calls to functions marked as returning Never / NoReturn (#18333)

## Summary

`ty` does not understand that calls to functions which have been
annotated as having a return type of `Never` / `NoReturn` are terminal.

This PR fixes that, by adding new reachability constraints when call
expressions are seen. If the call expression evaluates to `Never`, the
code following it will be considered to be unreachable. Note that, for
adding these constraints, we only consider call expressions at the
statement level, and that too only inside function scopes. This is
because otherwise, the number of such constraints becomes too high, and
evaluating them later on during type inference results in a major
performance degradation.

Fixes https://github.com/astral-sh/ty/issues/180

## Test Plan

New mdtests.

## Ecosystem changes

This PR removes the following false-positives:
- "Function can implicitly return `None`, which is not assignable to
...".
- "Name `foo` used when possibly not defind" - because the branch in
which it is not defined has a `NoReturn` call, or when `foo` was
imported in a `try`, and the except had a `NoReturn` call.

---------

Co-authored-by: David Peter <mail@david-peter.de>

crates/ty_python_semantic/src/semantic_index/builder.rs

@@ -35,8 +35,8 @@ use crate::semantic_index::place::{
     PlaceExprWithFlags, PlaceTableBuilder, Scope, ScopeId, ScopeKind, ScopedPlaceId,
 };
 use crate::semantic_index::predicate::{
-    PatternPredicate, PatternPredicateKind, Predicate, PredicateNode, PredicateOrLiteral,
-    ScopedPredicateId, StarImportPlaceholderPredicate,
+    CallableAndCallExpr, PatternPredicate, PatternPredicateKind, Predicate, PredicateNode,
+    PredicateOrLiteral, ScopedPredicateId, StarImportPlaceholderPredicate,
 };
 use crate::semantic_index::re_exports::exported_names;
 use crate::semantic_index::reachability_constraints::{
@@ -1901,11 +1901,45 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> {
                 value,
                 range: _,
                 node_index: _,
-            }) if self.in_module_scope() => {
-                if let Some(expr) = dunder_all_extend_argument(value) {
-                    self.add_standalone_expression(expr);
+            }) => {
+                if self.in_module_scope() {
+                    if let Some(expr) = dunder_all_extend_argument(value) {
+                        self.add_standalone_expression(expr);
+                    }
                 }
+
                 self.visit_expr(value);
+
+                // If the statement is a call, it could possibly be a call to a function
+                // marked with `NoReturn` (for example, `sys.exit()`). In this case, we use a special
+                // kind of constraint to mark the following code as unreachable.
+                //
+                // Ideally, these constraints should be added for every call expression, even those in
+                // sub-expressions and in the module-level scope. But doing so makes the number of
+                // such constraints so high that it significantly degrades performance. We thus cut
+                // scope here and add these constraints only at statement level function calls,
+                // like `sys.exit()`, and not within sub-expression like `3 + sys.exit()` etc.
+                //
+                // We also only add these inside function scopes, since considering module-level
+                // constraints can affect the the type of imported symbols, leading to a lot more
+                // work in third-party code.
+                if let ast::Expr::Call(ast::ExprCall { func, .. }) = value.as_ref() {
+                    if !self.source_type.is_stub() && self.in_function_scope() {
+                        let callable = self.add_standalone_expression(func);
+                        let call_expr = self.add_standalone_expression(value.as_ref());
+
+                        let predicate = Predicate {
+                            node: PredicateNode::ReturnsNever(CallableAndCallExpr {
+                                callable,
+                                call_expr,
+                            }),
+                            is_positive: false,
+                        };
+                        self.record_reachability_constraint(PredicateOrLiteral::Predicate(
+                            predicate,
+                        ));
+                    }
+                }
             }
             _ => {
                 walk_stmt(self, stmt);

crates/ty_python_semantic/src/semantic_index/predicate.rs

@@ -102,9 +102,16 @@ impl PredicateOrLiteral<'_> {
     }
 }
 
+#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq, salsa::Update, get_size2::GetSize)]
+pub(crate) struct CallableAndCallExpr<'db> {
+    pub(crate) callable: Expression<'db>,
+    pub(crate) call_expr: Expression<'db>,
+}
+
 #[derive(Clone, Copy, Debug, Hash, PartialEq, Eq, salsa::Update, get_size2::GetSize)]
 pub(crate) enum PredicateNode<'db> {
     Expression(Expression<'db>),
+    ReturnsNever(CallableAndCallExpr<'db>),
     Pattern(PatternPredicate<'db>),
     StarImportPlaceholder(StarImportPlaceholderPredicate<'db>),
 }

crates/ty_python_semantic/src/semantic_index/reachability_constraints.rs

@@ -204,7 +204,8 @@ use crate::place::{RequiresExplicitReExport, imported_symbol};
 use crate::semantic_index::expression::Expression;
 use crate::semantic_index::place_table;
 use crate::semantic_index::predicate::{
-    PatternPredicate, PatternPredicateKind, Predicate, PredicateNode, Predicates, ScopedPredicateId,
+    CallableAndCallExpr, PatternPredicate, PatternPredicateKind, Predicate, PredicateNode,
+    Predicates, ScopedPredicateId,
 };
 use crate::types::{Truthiness, Type, infer_expression_type};
 
@@ -684,6 +685,53 @@ impl ReachabilityConstraints {
                 let ty = infer_expression_type(db, test_expr);
                 ty.bool(db).negate_if(!predicate.is_positive)
             }
+            PredicateNode::ReturnsNever(CallableAndCallExpr {
+                callable,
+                call_expr,
+            }) => {
+                // We first infer just the type of the callable. In the most likely case that the
+                // function is not marked with `NoReturn`, or that it always returns `NoReturn`,
+                // doing so allows us to avoid the more expensive work of inferring the entire call
+                // expression (which could involve inferring argument types to possibly run the overload
+                // selection algorithm).
+                // Avoiding this on the happy-path is important because these constraints can be
+                // very large in number, since we add them on all statement level function calls.
+                let ty = infer_expression_type(db, callable);
+
+                let overloads_iterator =
+                    if let Some(Type::Callable(callable)) = ty.into_callable(db) {
+                        callable.signatures(db).overloads.iter()
+                    } else {
+                        return Truthiness::AlwaysFalse.negate_if(!predicate.is_positive);
+                    };
+
+                let (no_overloads_return_never, all_overloads_return_never) = overloads_iterator
+                    .fold((true, true), |(none, all), overload| {
+                        let overload_returns_never =
+                            overload.return_ty.is_some_and(|return_type| {
+                                return_type.is_equivalent_to(db, Type::Never)
+                            });
+
+                        (
+                            none && !overload_returns_never,
+                            all && overload_returns_never,
+                        )
+                    });
+
+                if no_overloads_return_never {
+                    Truthiness::AlwaysFalse
+                } else if all_overloads_return_never {
+                    Truthiness::AlwaysTrue
+                } else {
+                    let call_expr_ty = infer_expression_type(db, call_expr);
+                    if call_expr_ty.is_equivalent_to(db, Type::Never) {
+                        Truthiness::AlwaysTrue
+                    } else {
+                        Truthiness::AlwaysFalse
+                    }
+                }
+                .negate_if(!predicate.is_positive)
+            }
             PredicateNode::Pattern(inner) => Self::analyze_single_pattern_predicate(db, inner),
             PredicateNode::StarImportPlaceholder(star_import) => {
                 let place_table = place_table(db, star_import.scope(db));

crates/ty_python_semantic/src/semantic_index/use_def.rs

@@ -192,6 +192,17 @@
 //! for that place that we need for that use or definition. When we reach the end of the scope, it
 //! records the state for each place as the public definitions of that place.
 //!
+//! ```python
+//! x = 1
+//! x = 2
+//! y = x
+//! if flag:
+//!     x = 3
+//! else:
+//!     x = 4
+//! z = x
+//! ```
+//!
 //! Let's walk through the above example. Initially we do not have any record of `x`. When we add
 //! the new place (before we process the first binding), we create a new undefined `PlaceState`
 //! which has a single live binding (the "unbound" definition) and a single live declaration (the

crates/ty_python_semantic/src/types.rs

@@ -7226,7 +7226,7 @@ impl<'db> BoundMethodType<'db> {
 #[derive(PartialOrd, Ord)]
 pub struct CallableType<'db> {
     #[returns(ref)]
-    signatures: CallableSignature<'db>,
+    pub(crate) signatures: CallableSignature<'db>,
 
     /// We use `CallableType` to represent function-like objects, like the synthesized methods
     /// of dataclasses or NamedTuples. These callables act like real functions when accessed

crates/ty_python_semantic/src/types/infer.rs

@@ -2122,7 +2122,9 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                 node_index: _,
                 value,
             }) => {
-                self.infer_expression(value);
+                // If this is a call expression, we would have added a `ReturnsNever` constraint,
+                // meaning this will be a standalone expression.
+                self.infer_maybe_standalone_expression(value);
             }
             ast::Stmt::If(if_statement) => self.infer_if_statement(if_statement),
             ast::Stmt::Try(try_statement) => self.infer_try_statement(try_statement),
@@ -5263,7 +5265,8 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
         // arguments after matching them to parameters, but before checking that the argument types
         // are assignable to any parameter annotations.
         let call_arguments = Self::parse_arguments(arguments);
-        let callable_type = self.infer_expression(func);
+
+        let callable_type = self.infer_maybe_standalone_expression(func);
 
         if let Type::FunctionLiteral(function) = callable_type {
             // Make sure that the `function.definition` is only called when the function is defined

crates/ty_python_semantic/src/types/narrow.rs

@@ -3,7 +3,7 @@ use crate::semantic_index::expression::Expression;
 use crate::semantic_index::place::{PlaceExpr, PlaceTable, ScopeId, ScopedPlaceId};
 use crate::semantic_index::place_table;
 use crate::semantic_index::predicate::{
-    PatternPredicate, PatternPredicateKind, Predicate, PredicateNode,
+    CallableAndCallExpr, PatternPredicate, PatternPredicateKind, Predicate, PredicateNode,
 };
 use crate::types::function::KnownFunction;
 use crate::types::infer::infer_same_file_expression_type;
@@ -59,6 +59,7 @@ pub(crate) fn infer_narrowing_constraint<'db>(
                 all_negative_narrowing_constraints_for_pattern(db, pattern)
             }
         }
+        PredicateNode::ReturnsNever(_) => return None,
         PredicateNode::StarImportPlaceholder(_) => return None,
     };
     if let Some(constraints) = constraints {
@@ -346,6 +347,7 @@ impl<'db, 'ast> NarrowingConstraintsBuilder<'db, 'ast> {
             PredicateNode::Pattern(pattern) => {
                 self.evaluate_pattern_predicate(pattern, self.is_positive)
             }
+            PredicateNode::ReturnsNever(_) => return None,
             PredicateNode::StarImportPlaceholder(_) => return None,
         };
         if let Some(mut constraints) = constraints {
@@ -429,6 +431,9 @@ impl<'db, 'ast> NarrowingConstraintsBuilder<'db, 'ast> {
         match self.predicate {
             PredicateNode::Expression(expression) => expression.scope(self.db),
             PredicateNode::Pattern(pattern) => pattern.scope(self.db),
+            PredicateNode::ReturnsNever(CallableAndCallExpr { callable, .. }) => {
+                callable.scope(self.db)
+            }
             PredicateNode::StarImportPlaceholder(definition) => definition.scope(self.db),
         }
     }