Skip to content

Commit

Permalink
Add autotyping-like return type inference for annotation rules (#8643)
Browse files Browse the repository at this point in the history
## Summary

This PR adds (unsafe) fixes to the flake8-annotations rules that enforce
missing return types, offering to automatically insert type annotations
for functions with literal return values. The logic is smart enough to
generate simplified unions (e.g., `float` instead of `int | float`) and
deal with implicit returns (`return` without a value).

Closes #1640 (though we could
open a separate issue for referring parameter types).

Closes #8213.

## Test Plan

`cargo test`
  • Loading branch information
charliermarsh authored Nov 14, 2023
1 parent 23c819b commit bf2cc3f
Show file tree
Hide file tree
Showing 18 changed files with 580 additions and 140 deletions.
Original file line number Diff line number Diff line change
@@ -0,0 +1,32 @@
def func():
return 1


def func():
return 1.5


def func(x: int):
if x > 0:
return 1
else:
return 1.5


def func():
return True


def func(x: int):
if x > 0:
return None
else:
return


def func(x: int):
return 1 or 2.5 if x > 0 else 1.5 or "str"


def func(x: int):
return 1 + 2.5 if x > 0 else 1.5 or "str"
87 changes: 87 additions & 0 deletions crates/ruff_linter/src/rules/flake8_annotations/helpers.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,14 @@
use itertools::Itertools;

use ruff_python_ast::helpers::{pep_604_union, ReturnStatementVisitor};
use ruff_python_ast::visitor::Visitor;
use ruff_python_ast::{self as ast, Expr, ExprContext};
use ruff_python_semantic::analyze::type_inference::{NumberLike, PythonType, ResolvedPythonType};
use ruff_python_semantic::analyze::visibility;
use ruff_python_semantic::{Definition, SemanticModel};
use ruff_text_size::TextRange;

use crate::settings::types::PythonVersion;

/// Return the name of the function, if it's overloaded.
pub(crate) fn overloaded_name(definition: &Definition, semantic: &SemanticModel) -> Option<String> {
Expand Down Expand Up @@ -27,3 +36,81 @@ pub(crate) fn is_overload_impl(
function.name.as_str() == overloaded_name
}
}

/// Given a function, guess its return type.
pub(crate) fn auto_return_type(
function: &ast::StmtFunctionDef,
target_version: PythonVersion,
) -> Option<Expr> {
// Collect all the `return` statements.
let returns = {
let mut visitor = ReturnStatementVisitor::default();
visitor.visit_body(&function.body);
if visitor.is_generator {
return None;
}
visitor.returns
};

// Determine the return type of the first `return` statement.
let (return_statement, returns) = returns.split_first()?;
let mut return_type = return_statement.value.as_deref().map_or(
ResolvedPythonType::Atom(PythonType::None),
ResolvedPythonType::from,
);

// Merge the return types of the remaining `return` statements.
for return_statement in returns {
return_type = return_type.union(return_statement.value.as_deref().map_or(
ResolvedPythonType::Atom(PythonType::None),
ResolvedPythonType::from,
));
}

match return_type {
ResolvedPythonType::Atom(python_type) => type_expr(python_type),
ResolvedPythonType::Union(python_types) if target_version >= PythonVersion::Py310 => {
// Aggregate all the individual types (e.g., `int`, `float`).
let names = python_types
.iter()
.sorted_unstable()
.filter_map(|python_type| type_expr(*python_type))
.collect::<Vec<_>>();

// Wrap in a bitwise union (e.g., `int | float`).
Some(pep_604_union(&names))
}
ResolvedPythonType::Union(_) => None,
ResolvedPythonType::Unknown => None,
ResolvedPythonType::TypeError => None,
}
}

/// Given a [`PythonType`], return an [`Expr`] that resolves to that type.
fn type_expr(python_type: PythonType) -> Option<Expr> {
fn name(name: &str) -> Expr {
Expr::Name(ast::ExprName {
id: name.into(),
range: TextRange::default(),
ctx: ExprContext::Load,
})
}

match python_type {
PythonType::String => Some(name("str")),
PythonType::Bytes => Some(name("bytes")),
PythonType::Number(number) => match number {
NumberLike::Integer => Some(name("int")),
NumberLike::Float => Some(name("float")),
NumberLike::Complex => Some(name("complex")),
NumberLike::Bool => Some(name("bool")),
},
PythonType::None => Some(name("None")),
PythonType::Ellipsis => None,
PythonType::Dict => None,
PythonType::List => None,
PythonType::Set => None,
PythonType::Tuple => None,
PythonType::Generator => None,
}
}
18 changes: 18 additions & 0 deletions crates/ruff_linter/src/rules/flake8_annotations/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -110,6 +110,24 @@ mod tests {
Ok(())
}

#[test]
fn auto_return_type() -> Result<()> {
let diagnostics = test_path(
Path::new("flake8_annotations/auto_return_type.py"),
&LinterSettings {
..LinterSettings::for_rules(vec![
Rule::MissingReturnTypeUndocumentedPublicFunction,
Rule::MissingReturnTypePrivateFunction,
Rule::MissingReturnTypeSpecialMethod,
Rule::MissingReturnTypeStaticMethod,
Rule::MissingReturnTypeClassMethod,
])
},
)?;
assert_messages!(diagnostics);
Ok(())
}

#[test]
fn suppress_none_returning() -> Result<()> {
let diagnostics = test_path(
Expand Down
Loading

0 comments on commit bf2cc3f

Please sign in to comment.