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 #9201.
astral-sh · Dec 20, 2023 · 5ccc21a · 5ccc21a
1 parent f5d4019
commit 5ccc21a
Show file tree

Hide file tree

Showing 5 changed files with 406 additions and 150 deletions.
diff --git a/crates/ruff_linter/resources/test/fixtures/flake8_annotations/auto_return_type.py b/crates/ruff_linter/resources/test/fixtures/flake8_annotations/auto_return_type.py
@@ -182,3 +182,33 @@ def method(self):
             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
diff --git a/crates/ruff_linter/src/rules/flake8_annotations/helpers.rs b/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![]))

diff --git a/...nnotations/snapshots/ruff_linter__rules__flake8_annotations__tests__auto_return_type.snap b/...nnotations/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:
+
 
diff --git a/...tions/snapshots/ruff_linter__rules__flake8_annotations__tests__auto_return_type_py38.snap b/...tions/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:
+