Respect mixed return and raise cases in return-type analysis (#9310)

## Summary Given: ```python from somewhere import get_cfg def lookup_cfg(cfg_description): cfg = get_cfg(cfg_description) if cfg is not None: return cfg raise AttributeError(f"No cfg found matching {cfg_description}") ``` We were analyzing the method from last-to-first statement. So we saw the `raise`, then assumed the method _always_ raised. In reality, though, it _might_ return. This PR improves the branch analysis to respect these mixed cases. Closes #9269. Closes #9304.
astral-sh · Dec 29, 2023 · 2895e7d · 2895e7d
1 parent 00f3c7d
commit 2895e7d
Show file tree

Hide file tree

Showing 7 changed files with 482 additions and 203 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
@@ -229,3 +229,38 @@ def overloaded(i: "str") -> "str":
 
 def overloaded(i):
     return i
+
+
+def func(x: int):
+    if not x:
+        return 1
+    raise ValueError
+
+
+def func(x: int):
+    if not x:
+        return 1
+    else:
+        return 2
+    raise ValueError
+
+
+def func():
+    try:
+        raise ValueError
+    except:
+        return 2
+
+
+def func():
+    try:
+        return 1
+    except:
+        pass
+
+
+def func(x: int):
+    for _ in range(3):
+        if x > 0:
+            return 1
+    raise ValueError
diff --git a/crates/ruff_linter/src/rules/flake8_annotations/helpers.rs b/crates/ruff_linter/src/rules/flake8_annotations/helpers.rs
@@ -3,10 +3,11 @@ use rustc_hash::FxHashSet;
 
 use ruff_diagnostics::Edit;
 use ruff_python_ast::helpers::{
-    pep_604_union, typing_optional, typing_union, ReturnStatementVisitor, Terminal,
+    pep_604_union, typing_optional, typing_union, ReturnStatementVisitor,
 };
 use ruff_python_ast::visitor::Visitor;
 use ruff_python_ast::{self as ast, Expr, ExprContext};
+use ruff_python_semantic::analyze::terminal::Terminal;
 use ruff_python_semantic::analyze::type_inference::{NumberLike, PythonType, ResolvedPythonType};
 use ruff_python_semantic::analyze::visibility;
 use ruff_python_semantic::{Definition, SemanticModel};
@@ -61,7 +62,7 @@ pub(crate) fn auto_return_type(function: &ast::StmtFunctionDef) -> Option<AutoPy
     let terminal = Terminal::from_function(function);
 
     // If every control flow path raises an exception, return `NoReturn`.
-    if terminal == Some(Terminal::Raise) {
+    if terminal == Terminal::Raise {
         return Some(AutoPythonType::Never);
     }
 
@@ -88,7 +89,7 @@ pub(crate) fn auto_return_type(function: &ast::StmtFunctionDef) -> Option<AutoPy
     //     if x > 0:
     //         return 1
     // ```
-    if terminal.is_none() {
+    if terminal == Terminal::ConditionalReturn || terminal == Terminal::None {
         return_type = return_type.union(ResolvedPythonType::Atom(PythonType::None));
     }
 

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
@@ -579,4 +579,99 @@ auto_return_type.py:210:5: ANN201 [*] Missing return type annotation for public
 212 212 |         raise ValueError
 213 213 |     else:
 
+auto_return_type.py:234:5: ANN201 [*] Missing return type annotation for public function `func`
+    |
+234 | def func(x: int):
+    |     ^^^^ ANN201
+235 |     if not x:
+236 |         return 1
+    |
+    = help: Add return type annotation: `int`
+
+ℹ Unsafe fix
+231 231 |     return i
+232 232 | 
+233 233 | 
+234     |-def func(x: int):
+    234 |+def func(x: int) -> int:
+235 235 |     if not x:
+236 236 |         return 1
+237 237 |     raise ValueError
+
+auto_return_type.py:240:5: ANN201 [*] Missing return type annotation for public function `func`
+    |
+240 | def func(x: int):
+    |     ^^^^ ANN201
+241 |     if not x:
+242 |         return 1
+    |
+    = help: Add return type annotation: `int`
+
+ℹ Unsafe fix
+237 237 |     raise ValueError
+238 238 | 
+239 239 | 
+240     |-def func(x: int):
+    240 |+def func(x: int) -> int:
+241 241 |     if not x:
+242 242 |         return 1
+243 243 |     else:
+
+auto_return_type.py:248:5: ANN201 [*] Missing return type annotation for public function `func`
+    |
+248 | def func():
+    |     ^^^^ ANN201
+249 |     try:
+250 |         raise ValueError
+    |
+    = help: Add return type annotation: `int | None`
+
+ℹ Unsafe fix
+245 245 |     raise ValueError
+246 246 | 
+247 247 | 
+248     |-def func():
+    248 |+def func() -> int | None:
+249 249 |     try:
+250 250 |         raise ValueError
+251 251 |     except:
+
+auto_return_type.py:255:5: ANN201 [*] Missing return type annotation for public function `func`
+    |
+255 | def func():
+    |     ^^^^ ANN201
+256 |     try:
+257 |         return 1
+    |
+    = help: Add return type annotation: `int | None`
+
+ℹ Unsafe fix
+252 252 |         return 2
+253 253 | 
+254 254 | 
+255     |-def func():
+    255 |+def func() -> int | None:
+256 256 |     try:
+257 257 |         return 1
+258 258 |     except:
+
+auto_return_type.py:262:5: ANN201 [*] Missing return type annotation for public function `func`
+    |
+262 | def func(x: int):
+    |     ^^^^ ANN201
+263 |     for _ in range(3):
+264 |         if x > 0:
+    |
+    = help: Add return type annotation: `int`
+
+ℹ Unsafe fix
+259 259 |         pass
+260 260 | 
+261 261 | 
+262     |-def func(x: int):
+    262 |+def func(x: int) -> int:
+263 263 |     for _ in range(3):
+264 264 |         if x > 0:
+265 265 |             return 1
+
 
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
@@ -642,4 +642,117 @@ auto_return_type.py:210:5: ANN201 [*] Missing return type annotation for public
 212 212 |         raise ValueError
 213 213 |     else:
 
+auto_return_type.py:234:5: ANN201 [*] Missing return type annotation for public function `func`
+    |
+234 | def func(x: int):
+    |     ^^^^ ANN201
+235 |     if not x:
+236 |         return 1
+    |
+    = help: Add return type annotation: `int`
+
+ℹ Unsafe fix
+231 231 |     return i
+232 232 | 
+233 233 | 
+234     |-def func(x: int):
+    234 |+def func(x: int) -> int:
+235 235 |     if not x:
+236 236 |         return 1
+237 237 |     raise ValueError
+
+auto_return_type.py:240:5: ANN201 [*] Missing return type annotation for public function `func`
+    |
+240 | def func(x: int):
+    |     ^^^^ ANN201
+241 |     if not x:
+242 |         return 1
+    |
+    = help: Add return type annotation: `int`
+
+ℹ Unsafe fix
+237 237 |     raise ValueError
+238 238 | 
+239 239 | 
+240     |-def func(x: int):
+    240 |+def func(x: int) -> int:
+241 241 |     if not x:
+242 242 |         return 1
+243 243 |     else:
+
+auto_return_type.py:248:5: ANN201 [*] Missing return type annotation for public function `func`
+    |
+248 | def func():
+    |     ^^^^ ANN201
+249 |     try:
+250 |         raise ValueError
+    |
+    = help: Add return type annotation: `Optional[int]`
+
+ℹ Unsafe fix
+214 214 |         return 1
+215 215 | 
+216 216 | 
+217     |-from typing import overload
+    217 |+from typing import overload, Optional
+218 218 | 
+219 219 | 
+220 220 | @overload
+--------------------------------------------------------------------------------
+245 245 |     raise ValueError
+246 246 | 
+247 247 | 
+248     |-def func():
+    248 |+def func() -> Optional[int]:
+249 249 |     try:
+250 250 |         raise ValueError
+251 251 |     except:
+
+auto_return_type.py:255:5: ANN201 [*] Missing return type annotation for public function `func`
+    |
+255 | def func():
+    |     ^^^^ ANN201
+256 |     try:
+257 |         return 1
+    |
+    = help: Add return type annotation: `Optional[int]`
+
+ℹ Unsafe fix
+214 214 |         return 1
+215 215 | 
+216 216 | 
+217     |-from typing import overload
+    217 |+from typing import overload, Optional
+218 218 | 
+219 219 | 
+220 220 | @overload
+--------------------------------------------------------------------------------
+252 252 |         return 2
+253 253 | 
+254 254 | 
+255     |-def func():
+    255 |+def func() -> Optional[int]:
+256 256 |     try:
+257 257 |         return 1
+258 258 |     except:
+
+auto_return_type.py:262:5: ANN201 [*] Missing return type annotation for public function `func`
+    |
+262 | def func(x: int):
+    |     ^^^^ ANN201
+263 |     for _ in range(3):
+264 |         if x > 0:
+    |
+    = help: Add return type annotation: `int`
+
+ℹ Unsafe fix
+259 259 |         pass
+260 260 | 
+261 261 | 
+262     |-def func(x: int):
+    262 |+def func(x: int) -> int:
+263 263 |     for _ in range(3):
+264 264 |         if x > 0:
+265 265 |             return 1
+