Add support for NoReturn in auto-return-typing (#9206)

## Summary

Given a function like:

```python
def func(x: int):
    if not x:
        raise ValueError
    else:
        raise TypeError
```

We now correctly use `NoReturn` as the return type, rather than `None`.

Closes https://github.com/astral-sh/ruff/issues/9201.

crates/ruff_linter/resources/test/fixtures/flake8_annotations/auto_return_type.py

@@ -182,3 +182,33 @@ class Foo(abc.ABC):
             return 1
         else:
             return 1.5
+
+
+def func(x: int):
+    try:
+        pass
+    except:
+        return 2
+
+
+def func(x: int):
+    try:
+        pass
+    except:
+        return 2
+    else:
+        return 3
+
+
+def func(x: int):
+    if not x:
+        raise ValueError
+    else:
+        raise TypeError
+
+
+def func(x: int):
+    if not x:
+        raise ValueError
+    else:
+        return 1

crates/ruff_linter/src/rules/flake8_annotations/helpers.rs

@@ -3,7 +3,7 @@ use rustc_hash::FxHashSet;
 
 use ruff_diagnostics::Edit;
 use ruff_python_ast::helpers::{
-    implicit_return, pep_604_union, typing_optional, typing_union, ReturnStatementVisitor,
+    pep_604_union, typing_optional, typing_union, ReturnStatementVisitor, Terminal,
 };
 use ruff_python_ast::visitor::Visitor;
 use ruff_python_ast::{self as ast, Expr, ExprContext};
@@ -57,6 +57,14 @@ pub(crate) fn auto_return_type(function: &ast::StmtFunctionDef) -> Option<AutoPy
         visitor.returns
     };
 
+    // Determine the terminal behavior (i.e., implicit return, no return, etc.).
+    let terminal = Terminal::from_function(function);
+
+    // If every control flow path raises an exception, return `NoReturn`.
+    if terminal == Some(Terminal::Raise) {
+        return Some(AutoPythonType::NoReturn);
+    }
+
     // Determine the return type of the first `return` statement.
     let Some((return_statement, returns)) = returns.split_first() else {
         return Some(AutoPythonType::Atom(PythonType::None));
@@ -80,7 +88,7 @@ pub(crate) fn auto_return_type(function: &ast::StmtFunctionDef) -> Option<AutoPy
     //     if x > 0:
     //         return 1
     // ```
-    if implicit_return(function) {
+    if terminal.is_none() {
         return_type = return_type.union(ResolvedPythonType::Atom(PythonType::None));
     }
 
@@ -94,6 +102,7 @@ pub(crate) fn auto_return_type(function: &ast::StmtFunctionDef) -> Option<AutoPy
 
 #[derive(Debug)]
 pub(crate) enum AutoPythonType {
+    NoReturn,
     Atom(PythonType),
     Union(FxHashSet<PythonType>),
 }
@@ -111,6 +120,21 @@ impl AutoPythonType {
         target_version: PythonVersion,
     ) -> Option<(Expr, Vec<Edit>)> {
         match self {
+            AutoPythonType::NoReturn => {
+                let (no_return_edit, binding) = importer
+                    .get_or_import_symbol(
+                        &ImportRequest::import_from("typing", "NoReturn"),
+                        at,
+                        semantic,
+                    )
+                    .ok()?;
+                let expr = Expr::Name(ast::ExprName {
+                    id: binding,
+                    range: TextRange::default(),
+                    ctx: ExprContext::Load,
+                });
+                Some((expr, vec![no_return_edit]))
+            }
             AutoPythonType::Atom(python_type) => {
                 let expr = type_expr(python_type)?;
                 Some((expr, vec![]))

crates/ruff_linter/src/rules/flake8_annotations/snapshots/ruff_linter__rules__flake8_annotations__tests__auto_return_type.snap

@@ -495,4 +495,88 @@ auto_return_type.py:180:9: ANN201 [*] Missing return type annotation for public
 182 182 |             return 1
 183 183 |         else:
 
+auto_return_type.py:187:5: ANN201 [*] Missing return type annotation for public function `func`
+    |
+187 | def func(x: int):
+    |     ^^^^ ANN201
+188 |     try:
+189 |         pass
+    |
+    = help: Add return type annotation: `int | None`
+
+ℹ Unsafe fix
+184 184 |             return 1.5
+185 185 | 
+186 186 | 
+187     |-def func(x: int):
+    187 |+def func(x: int) -> int | None:
+188 188 |     try:
+189 189 |         pass
+190 190 |     except:
+
+auto_return_type.py:194:5: ANN201 [*] Missing return type annotation for public function `func`
+    |
+194 | def func(x: int):
+    |     ^^^^ ANN201
+195 |     try:
+196 |         pass
+    |
+    = help: Add return type annotation: `int`
+
+ℹ Unsafe fix
+191 191 |         return 2
+192 192 | 
+193 193 | 
+194     |-def func(x: int):
+    194 |+def func(x: int) -> int:
+195 195 |     try:
+196 196 |         pass
+197 197 |     except:
+
+auto_return_type.py:203:5: ANN201 [*] Missing return type annotation for public function `func`
+    |
+203 | def func(x: int):
+    |     ^^^^ ANN201
+204 |     if not x:
+205 |         raise ValueError
+    |
+    = help: Add return type annotation: `NoReturn`
+
+ℹ Unsafe fix
+151 151 | 
+152 152 | import abc
+153 153 | from abc import abstractmethod
+    154 |+from typing import NoReturn
+154 155 | 
+155 156 | 
+156 157 | class Foo(abc.ABC):
+--------------------------------------------------------------------------------
+200 201 |         return 3
+201 202 | 
+202 203 | 
+203     |-def func(x: int):
+    204 |+def func(x: int) -> NoReturn:
+204 205 |     if not x:
+205 206 |         raise ValueError
+206 207 |     else:
+
+auto_return_type.py:210:5: ANN201 [*] Missing return type annotation for public function `func`
+    |
+210 | def func(x: int):
+    |     ^^^^ ANN201
+211 |     if not x:
+212 |         raise ValueError
+    |
+    = help: Add return type annotation: `int`
+
+ℹ Unsafe fix
+207 207 |         raise TypeError
+208 208 | 
+209 209 | 
+210     |-def func(x: int):
+    210 |+def func(x: int) -> int:
+211 211 |     if not x:
+212 212 |         raise ValueError
+213 213 |     else:
+
 

crates/ruff_linter/src/rules/flake8_annotations/snapshots/ruff_linter__rules__flake8_annotations__tests__auto_return_type_py38.snap

@@ -550,4 +550,96 @@ auto_return_type.py:180:9: ANN201 [*] Missing return type annotation for public
 182 182 |             return 1
 183 183 |         else:
 
+auto_return_type.py:187:5: ANN201 [*] Missing return type annotation for public function `func`
+    |
+187 | def func(x: int):
+    |     ^^^^ ANN201
+188 |     try:
+189 |         pass
+    |
+    = help: Add return type annotation: `Optional[int]`
+
+ℹ Unsafe fix
+151 151 | 
+152 152 | import abc
+153 153 | from abc import abstractmethod
+    154 |+from typing import Optional
+154 155 | 
+155 156 | 
+156 157 | class Foo(abc.ABC):
+--------------------------------------------------------------------------------
+184 185 |             return 1.5
+185 186 | 
+186 187 | 
+187     |-def func(x: int):
+    188 |+def func(x: int) -> Optional[int]:
+188 189 |     try:
+189 190 |         pass
+190 191 |     except:
+
+auto_return_type.py:194:5: ANN201 [*] Missing return type annotation for public function `func`
+    |
+194 | def func(x: int):
+    |     ^^^^ ANN201
+195 |     try:
+196 |         pass
+    |
+    = help: Add return type annotation: `int`
+
+ℹ Unsafe fix
+191 191 |         return 2
+192 192 | 
+193 193 | 
+194     |-def func(x: int):
+    194 |+def func(x: int) -> int:
+195 195 |     try:
+196 196 |         pass
+197 197 |     except:
+
+auto_return_type.py:203:5: ANN201 [*] Missing return type annotation for public function `func`
+    |
+203 | def func(x: int):
+    |     ^^^^ ANN201
+204 |     if not x:
+205 |         raise ValueError
+    |
+    = help: Add return type annotation: `NoReturn`
+
+ℹ Unsafe fix
+151 151 | 
+152 152 | import abc
+153 153 | from abc import abstractmethod
+    154 |+from typing import NoReturn
+154 155 | 
+155 156 | 
+156 157 | class Foo(abc.ABC):
+--------------------------------------------------------------------------------
+200 201 |         return 3
+201 202 | 
+202 203 | 
+203     |-def func(x: int):
+    204 |+def func(x: int) -> NoReturn:
+204 205 |     if not x:
+205 206 |         raise ValueError
+206 207 |     else:
+
+auto_return_type.py:210:5: ANN201 [*] Missing return type annotation for public function `func`
+    |
+210 | def func(x: int):
+    |     ^^^^ ANN201
+211 |     if not x:
+212 |         raise ValueError
+    |
+    = help: Add return type annotation: `int`
+
+ℹ Unsafe fix
+207 207 |         raise TypeError
+208 208 | 
+209 209 | 
+210     |-def func(x: int):
+    210 |+def func(x: int) -> int:
+211 211 |     if not x:
+212 212 |         raise ValueError
+213 213 |     else:
+
 

crates/ruff_python_ast/src/helpers.rs

@@ -921,178 +921,204 @@ where
     }
 }
 
-/// Returns `true` if the function has an implicit return.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
-pub fn implicit_return(function: &ast::StmtFunctionDef) -> bool {
+pub enum Terminal {
-    /// Returns `true` if the body may break via a `break` statement.
+    /// Every path through the function ends with a `raise` statement.
-    fn sometimes_breaks(stmts: &[Stmt]) -> bool {
+    Raise,
-        for stmt in stmts {
+    /// Every path through the function ends with a `return` (or `raise`) statement.
-            match stmt {
+    Return,
-                Stmt::For(ast::StmtFor { body, orelse, .. }) => {
+}
-                    if returns(body) {
-                        return false;
-                    }
-                    if sometimes_breaks(orelse) {
-                        return true;
-                    }
-                }
-                Stmt::While(ast::StmtWhile { body, orelse, .. }) => {
-                    if returns(body) {
-                        return false;
-                    }
-                    if sometimes_breaks(orelse) {
-                        return true;
-                    }
-                }
-                Stmt::If(ast::StmtIf {
-                    body,
-                    elif_else_clauses,
-                    ..
-                }) => {
-                    if std::iter::once(body)
-                        .chain(elif_else_clauses.iter().map(|clause| &clause.body))
-                        .any(|body| sometimes_breaks(body))
-                    {
-                        return true;
-                    }
-                }
-                Stmt::Match(ast::StmtMatch { cases, .. }) => {
-                    if cases.iter().any(|case| sometimes_breaks(&case.body)) {
-                        return true;
-                    }
-                }
-                Stmt::Try(ast::StmtTry {
-                    body,
-                    handlers,
-                    orelse,
-                    finalbody,
-                    ..
-                }) => {
-                    if sometimes_breaks(body)
-                        || handlers.iter().any(|handler| {
-                            let ExceptHandler::ExceptHandler(ast::ExceptHandlerExceptHandler {
-                                body,
-                                ..
-                            }) = handler;
-                            sometimes_breaks(body)
-                        })
-                        || sometimes_breaks(orelse)
-                        || sometimes_breaks(finalbody)
-                    {
-                        return true;
-                    }
-                }
-                Stmt::With(ast::StmtWith { body, .. }) => {
-                    if sometimes_breaks(body) {
-                        return true;
-                    }
-                }
-                Stmt::Break(_) => return true,
-                Stmt::Return(_) => return false,
-                Stmt::Raise(_) => return false,
-                _ => {}
-            }
-        }
-        false
-    }
 
-    /// Returns `true` if the body may break via a `break` statement.
+impl Terminal {
-    fn always_breaks(stmts: &[Stmt]) -> bool {
+    /// Returns the [`Terminal`] behavior of the function, if it can be determined, or `None` if the
-        for stmt in stmts {
+    /// function contains at least one control flow path that does not end with a `return` or `raise`
-            match stmt {
+    /// statement.
-                Stmt::Break(_) => return true,
+    pub fn from_function(function: &ast::StmtFunctionDef) -> Option<Terminal> {
-                Stmt::Return(_) => return false,
+        /// Returns `true` if the body may break via a `break` statement.
-                Stmt::Raise(_) => return false,
+        fn sometimes_breaks(stmts: &[Stmt]) -> bool {
-                _ => {}
+            for stmt in stmts {
-            }
+                match stmt {
-        }
+                    Stmt::For(ast::StmtFor { body, orelse, .. }) => {
-        false
+                        if returns(body).is_some() {
-    }
+                            return false;
-
+                        }
-    /// Returns `true` if the body contains a branch that ends without an explicit `return` or
+                        if sometimes_breaks(orelse) {
-    /// `raise` statement.
+                            return true;
-    fn returns(stmts: &[Stmt]) -> bool {
+                        }
-        for stmt in stmts.iter().rev() {
-            match stmt {
-                Stmt::For(ast::StmtFor { body, orelse, .. }) => {
-                    if always_breaks(body) {
-                        return false;
                     }
-                    if returns(body) {
+                    Stmt::While(ast::StmtWhile { body, orelse, .. }) => {
-                        return true;
+                        if returns(body).is_some() {
+                            return false;
+                        }
+                        if sometimes_breaks(orelse) {
+                            return true;
+                        }
                     }
-                    if returns(orelse) && !sometimes_breaks(body) {
+                    Stmt::If(ast::StmtIf {
-                        return true;
+                        body,
-                    }
+                        elif_else_clauses,
-                }
+                        ..
-                Stmt::While(ast::StmtWhile { body, orelse, .. }) => {
+                    }) => {
-                    if always_breaks(body) {
+                        if std::iter::once(body)
-                        return false;
-                    }
-                    if returns(body) {
-                        return true;
-                    }
-                    if returns(orelse) && !sometimes_breaks(body) {
-                        return true;
-                    }
-                }
-                Stmt::If(ast::StmtIf {
-                    body,
-                    elif_else_clauses,
-                    ..
-                }) => {
-                    if elif_else_clauses.iter().any(|clause| clause.test.is_none())
-                        && std::iter::once(body)
                             .chain(elif_else_clauses.iter().map(|clause| &clause.body))
-                            .all(|body| returns(body))
+                            .any(|body| sometimes_breaks(body))
-                    {
+                        {
-                        return true;
+                            return true;
+                        }
                     }
+                    Stmt::Match(ast::StmtMatch { cases, .. }) => {
+                        if cases.iter().any(|case| sometimes_breaks(&case.body)) {
+                            return true;
+                        }
+                    }
+                    Stmt::Try(ast::StmtTry {
+                        body,
+                        handlers,
+                        orelse,
+                        finalbody,
+                        ..
+                    }) => {
+                        if sometimes_breaks(body)
+                            || handlers.iter().any(|handler| {
+                                let ExceptHandler::ExceptHandler(ast::ExceptHandlerExceptHandler {
+                                    body,
+                                    ..
+                                }) = handler;
+                                sometimes_breaks(body)
+                            })
+                            || sometimes_breaks(orelse)
+                            || sometimes_breaks(finalbody)
+                        {
+                            return true;
+                        }
+                    }
+                    Stmt::With(ast::StmtWith { body, .. }) => {
+                        if sometimes_breaks(body) {
+                            return true;
+                        }
+                    }
+                    Stmt::Break(_) => return true,
+                    Stmt::Return(_) => return false,
+                    Stmt::Raise(_) => return false,
+                    _ => {}
                 }
-                Stmt::Match(ast::StmtMatch { cases, .. }) => {
-                    // Note: we assume the `match` is exhaustive.
-                    if cases.iter().all(|case| returns(&case.body)) {
-                        return true;
-                    }
-                }
-                Stmt::Try(ast::StmtTry {
-                    body,
-                    handlers,
-                    orelse,
-                    finalbody,
-                    ..
-                }) => {
-                    // If the `finally` block returns, the `try` block must also return.
-                    if returns(finalbody) {
-                        return true;
-                    }
-
-                    // If the `body` or the `else` block returns, the `try` block must also return.
-                    if (returns(body) || returns(orelse))
-                        && handlers.iter().all(|handler| {
-                            let ExceptHandler::ExceptHandler(ast::ExceptHandlerExceptHandler {
-                                body,
-                                ..
-                            }) = handler;
-                            returns(body)
-                        })
-                    {
-                        return true;
-                    }
-                }
-                Stmt::With(ast::StmtWith { body, .. }) => {
-                    if returns(body) {
-                        return true;
-                    }
-                }
-                Stmt::Return(_) => return true,
-                Stmt::Raise(_) => return true,
-                _ => {}
             }
+            false
         }
-        false
+
+        /// Returns `true` if the body may break via a `break` statement.
+        fn always_breaks(stmts: &[Stmt]) -> bool {
+            for stmt in stmts {
+                match stmt {
+                    Stmt::Break(_) => return true,
+                    Stmt::Return(_) => return false,
+                    Stmt::Raise(_) => return false,
+                    _ => {}
+                }
+            }
+            false
+        }
+
+        /// Returns `true` if the body contains a branch that ends without an explicit `return` or
+        /// `raise` statement.
+        fn returns(stmts: &[Stmt]) -> Option<Terminal> {
+            for stmt in stmts.iter().rev() {
+                match stmt {
+                    Stmt::For(ast::StmtFor { body, orelse, .. })
+                    | Stmt::While(ast::StmtWhile { body, orelse, .. }) => {
+                        if always_breaks(body) {
+                            return None;
+                        }
+                        if let Some(terminal) = returns(body) {
+                            return Some(terminal);
+                        }
+                        if !sometimes_breaks(body) {
+                            if let Some(terminal) = returns(orelse) {
+                                return Some(terminal);
+                            }
+                        }
+                    }
+                    Stmt::If(ast::StmtIf {
+                        body,
+                        elif_else_clauses,
+                        ..
+                    }) => {
+                        if elif_else_clauses.iter().any(|clause| clause.test.is_none()) {
+                            match Terminal::combine(std::iter::once(returns(body)).chain(
+                                elif_else_clauses.iter().map(|clause| returns(&clause.body)),
+                            )) {
+                                Some(Terminal::Raise) => return Some(Terminal::Raise),
+                                Some(Terminal::Return) => return Some(Terminal::Return),
+                                _ => {}
+                            }
+                        }
+                    }
+                    Stmt::Match(ast::StmtMatch { cases, .. }) => {
+                        // Note: we assume the `match` is exhaustive.
+                        match Terminal::combine(cases.iter().map(|case| returns(&case.body))) {
+                            Some(Terminal::Raise) => return Some(Terminal::Raise),
+                            Some(Terminal::Return) => return Some(Terminal::Return),
+                            _ => {}
+                        }
+                    }
+                    Stmt::Try(ast::StmtTry {
+                        body,
+                        handlers,
+                        orelse,
+                        finalbody,
+                        ..
+                    }) => {
+                        // If the `finally` block returns, the `try` block must also return.
+                        if let Some(terminal) = returns(finalbody) {
+                            return Some(terminal);
+                        }
+
+                        // If the body returns, the `try` block must also return.
+                        if returns(body) == Some(Terminal::Return) {
+                            return Some(Terminal::Return);
+                        }
+
+                        // If the else block and all the handlers return, the `try` block must also
+                        // return.
+                        if let Some(terminal) =
+                            Terminal::combine(std::iter::once(returns(orelse)).chain(
+                                handlers.iter().map(|handler| {
+                                    let ExceptHandler::ExceptHandler(
+                                        ast::ExceptHandlerExceptHandler { body, .. },
+                                    ) = handler;
+                                    returns(body)
+                                }),
+                            ))
+                        {
+                            return Some(terminal);
+                        }
+                    }
+                    Stmt::With(ast::StmtWith { body, .. }) => {
+                        if let Some(terminal) = returns(body) {
+                            return Some(terminal);
+                        }
+                    }
+                    Stmt::Return(_) => return Some(Terminal::Return),
+                    Stmt::Raise(_) => return Some(Terminal::Raise),
+                    _ => {}
+                }
+            }
+            None
+        }
+
+        returns(&function.body)
     }
 
-    !returns(&function.body)
+    /// Combine a series of [`Terminal`] operators.
+    fn combine(iter: impl Iterator<Item = Option<Terminal>>) -> Option<Terminal> {
+        iter.reduce(|acc, terminal| match (acc, terminal) {
+            (Some(Self::Raise), Some(Self::Raise)) => Some(Self::Raise),
+            (Some(_), Some(Self::Return)) => Some(Self::Return),
+            (Some(Self::Return), Some(_)) => Some(Self::Return),
+            _ => None,
+        })
+        .flatten()
+    }
 }
 
 /// A [`StatementVisitor`] that collects all `raise` statements in a function or method.