ANN201 doesn't handle pytest.raises context manager

[jogo-openai]

Description

ruff version: 0.9.4

code:

import pytest


def test_exception():
    with pytest.raises(ValueError):
        raise ValueError("This is a test exception")
    a = 1
    assert a == 1


def a_function():
    return 1


def no_return():
    pass

ruff check --select=ANN201 --diff --unsafe-fixes example.py

--- example.py
+++ example.py
@@ -1,16 +1,17 @@
 import pytest
+from typing import NoReturn


-def test_exception():
+def test_exception() -> NoReturn:
     with pytest.raises(ValueError):
         raise ValueError("This is a test exception")
     a = 1
     assert a == 1


-def a_function():
+def a_function() -> int:
     return 1


-def no_return():
+def no_return() -> None:
     pass

Would fix 3 errors.

After applying the fix mypy reports the error:

example.py:5: error: Implicit return in function which does not return  [misc]

[MichaReiser]

Thanks for the report. The issue is really only about

def test_exception():
    with pytest.raises(ValueError):
        raise ValueError("This is a test exception")
    a = 1
    assert a == 1

specifically that Ruff adds the return type Never which is incorrect because pytest.raises catches the exception and doesn't propagate it.

Since it is hard for ruff to know if a context manager handles an exception (because it can't resolve the context manager nor its methods if defined in another file). We have to explore whether we can provide any type suggestion in that case (by handling a context manager similar to try) or if we need to bail out of providing a fix.

Relevant code:

ruff/crates/ruff_python_semantic/src/analyze/terminal.rs

Lines 42 to 150 in 13ffb5b

	for stmt in stmts {
	match stmt {
	Stmt::For(ast::StmtFor { body, orelse, .. })
	| Stmt::While(ast::StmtWhile { body, orelse, .. }) => {
	if always_breaks(body) {
	continue;
	}
	terminal = terminal.and_then(Self::from_body(body));
	if !sometimes_breaks(body) {
	terminal = terminal.and_then(Self::from_body(orelse));
	}
	}
	Stmt::If(ast::StmtIf {
	body,
	elif_else_clauses,
	..
	}) => {
	let branch_terminal = Terminal::branches(
	std::iter::once(Self::from_body(body)).chain(
	elif_else_clauses
	.iter()
	.map(|clause| Self::from_body(&clause.body)),
	),
	);
	// If the `if` statement is known to be exhaustive (by way of including an
	// `else`)...
	if elif_else_clauses.iter().any(|clause| clause.test.is_none()) {
	// And all branches return, then the `if` statement returns.
	terminal = terminal.and_then(branch_terminal);
	} else if branch_terminal.has_any_return() {
	// Otherwise, if any branch returns, we know this can't be a
	// non-returning function.
	terminal = terminal.and_then(Terminal::ConditionalReturn);
	}
	}
	Stmt::Match(ast::StmtMatch { cases, .. }) => {
	let branch_terminal = terminal.and_then(Terminal::branches(
	cases.iter().map(|case| Self::from_body(&case.body)),
	));
	// If the `match` is known to be exhaustive (by way of including a wildcard
	// pattern)...
	if cases.iter().any(is_wildcard) {
	// And all branches return, then the `match` statement returns.
	terminal = terminal.and_then(branch_terminal);
	} else {
	// Otherwise, if any branch returns, we know this can't be a
	// non-returning function.
	if branch_terminal.has_any_return() {
	terminal = terminal.and_then(Terminal::ConditionalReturn);
	}
	}
	}
	Stmt::Try(ast::StmtTry {
	body,
	handlers,
	orelse,
	finalbody,
	..
	}) => {
	// If the body returns, then this can't be a non-returning function. We assume
	// that _any_ statement in the body could raise an exception, so we don't
	// consider the body to be exhaustive. In other words, we assume the exception
	// handlers exist for a reason.
	let body_terminal = Self::from_body(body);
	if body_terminal.has_any_return() {
	terminal = terminal.and_then(Terminal::ConditionalReturn);
	}
	// If the `finally` block returns, the `try` block must also return. (Similarly,
	// if the `finally` block raises, the `try` block must also raise.)
	terminal = terminal.and_then(Self::from_body(finalbody));
	let branch_terminal = Terminal::branches(handlers.iter().map(|handler| {
	let ExceptHandler::ExceptHandler(ast::ExceptHandlerExceptHandler {
	body,
	..
	}) = handler;
	Self::from_body(body)
	}));
	if orelse.is_empty() {
	// If there's no `else`, we may fall through, so only mark that this can't
	// be a non-returning function if any of the branches return.
	if branch_terminal.has_any_return() {
	terminal = terminal.and_then(Terminal::ConditionalReturn);
	}
	} else {
	// If there's an `else`, we won't fall through. If all the handlers and
	// the `else` block return,, the `try` block also returns.
	terminal =
	terminal.and_then(branch_terminal.branch(Terminal::from_body(orelse)));
	}
	}
	Stmt::With(ast::StmtWith { body, .. }) => {
	terminal = terminal.and_then(Self::from_body(body));
	}
	Stmt::Return(_) => {
	terminal = terminal.and_then(Terminal::RaiseOrReturn);
	}
	Stmt::Raise(_) => {
	terminal = terminal.and_then(Terminal::Raise);
	}
	_ => {}
	}
	}