Skip to content

Commit

Permalink
Enable auto-return-type involving Optional and Union annotations (#…
Browse files Browse the repository at this point in the history 
…8885)

## Summary

Previously, this was only supported for Python 3.10 and later, since we
always use the PEP 604-style unions.
  • Loading branch information
@charliermarsh
charliermarsh committed Nov 29, 2023
1 parent ec7456b commit 6435e4e
Show file tree
Hide file tree
Showing 10 changed files with 526 additions and 69 deletions.
18 changes: 17 additions & 1 deletion crates/ruff_linter/resources/test/fixtures/flake8_annotations/auto_return_type.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -42,8 +42,24 @@ def func(x: int):
return {"foo": 1}


def func():
def func(x: int):
if not x:
return 1
else:
return True


def func(x: int):
if not x:
return 1
else:
return None


def func(x: int):
if not x:
return 1
elif x > 5:
return "str"
else:
return None
110 changes: 91 additions & 19 deletions crates/ruff_linter/src/rules/flake8_annotations/helpers.rs
View file
Original file line number Diff line number Diff line change
@@ -1,12 +1,17 @@
use itertools::Itertools;
use ruff_diagnostics::Edit;
use rustc_hash::FxHashSet;

use ruff_python_ast::helpers::{pep_604_union, ReturnStatementVisitor};
use crate::importer::{ImportRequest, Importer};
use ruff_python_ast::helpers::{
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::type_inference::{NumberLike, PythonType, ResolvedPythonType};
use ruff_python_semantic::analyze::visibility;
use ruff_python_semantic::{Definition, SemanticModel};
use ruff_text_size::TextRange;
use ruff_text_size::{TextRange, TextSize};

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

Expand Down Expand Up @@ -38,10 +43,7 @@ pub(crate) fn is_overload_impl(
}

/// Given a function, guess its return type.
pub(crate) fn auto_return_type(
function: &ast::StmtFunctionDef,
target_version: PythonVersion,
) -> Option<Expr> {
pub(crate) fn auto_return_type(function: &ast::StmtFunctionDef) -> Option<AutoPythonType> {
// Collect all the `return` statements.
let returns = {
let mut visitor = ReturnStatementVisitor::default();
Expand All @@ -68,24 +70,94 @@ pub(crate) fn auto_return_type(
}

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()
.map(|python_type| type_expr(*python_type))
.collect::<Option<Vec<_>>>()?;

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

#[derive(Debug)]
pub(crate) enum AutoPythonType {
Atom(PythonType),
Union(FxHashSet<PythonType>),
}

impl AutoPythonType {
/// Convert an [`AutoPythonType`] into an [`Expr`].
///
/// If the [`Expr`] relies on importing any external symbols, those imports will be returned as
/// additional edits.
pub(crate) fn into_expression(
self,
importer: &Importer,
at: TextSize,
semantic: &SemanticModel,
target_version: PythonVersion,
) -> Option<(Expr, Vec<Edit>)> {
match self {
AutoPythonType::Atom(python_type) => {
let expr = type_expr(python_type)?;
Some((expr, vec![]))
}
AutoPythonType::Union(python_types) => {
if target_version >= PythonVersion::Py310 {
// Aggregate all the individual types (e.g., `int`, `float`).
let names = python_types
.iter()
.sorted_unstable()
.map(|python_type| type_expr(*python_type))
.collect::<Option<Vec<_>>>()?;

// Wrap in a bitwise union (e.g., `int | float`).
let expr = pep_604_union(&names);

Some((expr, vec![]))
} else {
let python_types = python_types
.into_iter()
.sorted_unstable()
.collect::<Vec<_>>();

match python_types.as_slice() {
[python_type, PythonType::None] | [PythonType::None, python_type] => {
let element = type_expr(*python_type)?;

// Ex) `Optional[int]`
let (optional_edit, binding) = importer
.get_or_import_symbol(
&ImportRequest::import_from("typing", "Optional"),
at,
semantic,
)
.ok()?;
let expr = typing_optional(element, binding);
Some((expr, vec![optional_edit]))
}
_ => {
let elements = python_types
.into_iter()
.map(type_expr)
.collect::<Option<Vec<_>>>()?;

// Ex) `Union[int, str]`
let (union_edit, binding) = importer
.get_or_import_symbol(
&ImportRequest::import_from("typing", "Union"),
at,
semantic,
)
.ok()?;
let expr = typing_union(&elements, binding);
Some((expr, vec![union_edit]))
}
}
}
}
}
}
}

/// Given a [`PythonType`], return an [`Expr`] that resolves to that type.
fn type_expr(python_type: PythonType) -> Option<Expr> {
fn name(name: &str) -> Expr {
Expand Down
20 changes: 20 additions & 0 deletions crates/ruff_linter/src/rules/flake8_annotations/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ mod tests {

use crate::assert_messages;
use crate::registry::Rule;
use crate::settings::types::PythonVersion;
use crate::settings::LinterSettings;
use crate::test::test_path;

Expand Down Expand Up @@ -128,6 +129,25 @@ mod tests {
Ok(())
}

#[test]
fn auto_return_type_py38() -> Result<()> {
let diagnostics = test_path(
Path::new("flake8_annotations/auto_return_type.py"),
&LinterSettings {
target_version: PythonVersion::Py38,
..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
116 changes: 80 additions & 36 deletions crates/ruff_linter/src/rules/flake8_annotations/rules/definition.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -725,39 +725,55 @@ pub(crate) fn definition(
) {
if is_method && visibility::is_classmethod(decorator_list, checker.semantic()) {
if checker.enabled(Rule::MissingReturnTypeClassMethod) {
let return_type = auto_return_type(function, checker.settings.target_version)
.map(|return_type| checker.generator().expr(&return_type));
let return_type = auto_return_type(function)
.and_then(|return_type| {
return_type.into_expression(
checker.importer(),
function.parameters.start(),
checker.semantic(),
checker.settings.target_version,
)
})
.map(|(return_type, edits)| (checker.generator().expr(&return_type), edits));
let mut diagnostic = Diagnostic::new(
MissingReturnTypeClassMethod {
name: name.to_string(),
annotation: return_type.clone(),
annotation: return_type.clone().map(|(return_type, ..)| return_type),
},
function.identifier(),
);
if let Some(return_type) = return_type {
diagnostic.set_fix(Fix::unsafe_edit(Edit::insertion(
format!(" -> {return_type}"),
function.parameters.range().end(),
)));
if let Some((return_type, edits)) = return_type {
diagnostic.set_fix(Fix::unsafe_edits(
Edit::insertion(format!(" -> {return_type}"), function.parameters.end()),
edits,
));
}
diagnostics.push(diagnostic);
}
} else if is_method && visibility::is_staticmethod(decorator_list, checker.semantic()) {
if checker.enabled(Rule::MissingReturnTypeStaticMethod) {
let return_type = auto_return_type(function, checker.settings.target_version)
.map(|return_type| checker.generator().expr(&return_type));
let return_type = auto_return_type(function)
.and_then(|return_type| {
return_type.into_expression(
checker.importer(),
function.parameters.start(),
checker.semantic(),
checker.settings.target_version,
)
})
.map(|(return_type, edits)| (checker.generator().expr(&return_type), edits));
let mut diagnostic = Diagnostic::new(
MissingReturnTypeStaticMethod {
name: name.to_string(),
annotation: return_type.clone(),
annotation: return_type.clone().map(|(return_type, ..)| return_type),
},
function.identifier(),
);
if let Some(return_type) = return_type {
diagnostic.set_fix(Fix::unsafe_edit(Edit::insertion(
format!(" -> {return_type}"),
function.parameters.range().end(),
)));
if let Some((return_type, edits)) = return_type {
diagnostic.set_fix(Fix::unsafe_edits(
Edit::insertion(format!(" -> {return_type}"), function.parameters.end()),
edits,
));
}
diagnostics.push(diagnostic);
}
Expand All @@ -775,7 +791,7 @@ pub(crate) fn definition(
);
diagnostic.set_fix(Fix::unsafe_edit(Edit::insertion(
" -> None".to_string(),
function.parameters.range().end(),
function.parameters.end(),
)));
diagnostics.push(diagnostic);
}
Expand All @@ -793,7 +809,7 @@ pub(crate) fn definition(
if let Some(return_type) = return_type {
diagnostic.set_fix(Fix::unsafe_edit(Edit::insertion(
format!(" -> {return_type}"),
function.parameters.range().end(),
function.parameters.end(),
)));
}
diagnostics.push(diagnostic);
Expand All @@ -802,42 +818,70 @@ pub(crate) fn definition(
match visibility {
visibility::Visibility::Public => {
if checker.enabled(Rule::MissingReturnTypeUndocumentedPublicFunction) {
let return_type =
auto_return_type(function, checker.settings.target_version)
.map(|return_type| checker.generator().expr(&return_type));
let return_type = auto_return_type(function)
.and_then(|return_type| {
return_type.into_expression(
checker.importer(),
function.parameters.start(),
checker.semantic(),
checker.settings.target_version,
)
})
.map(|(return_type, edits)| {
(checker.generator().expr(&return_type), edits)
});
let mut diagnostic = Diagnostic::new(
MissingReturnTypeUndocumentedPublicFunction {
name: name.to_string(),
annotation: return_type.clone(),
annotation: return_type
.clone()
.map(|(return_type, ..)| return_type),
},
function.identifier(),
);
if let Some(return_type) = return_type {
diagnostic.set_fix(Fix::unsafe_edit(Edit::insertion(
format!(" -> {return_type}"),
function.parameters.range().end(),
)));
if let Some((return_type, edits)) = return_type {
diagnostic.set_fix(Fix::unsafe_edits(
Edit::insertion(
format!(" -> {return_type}"),
function.parameters.end(),
),
edits,
));
}
diagnostics.push(diagnostic);
}
}
visibility::Visibility::Private => {
if checker.enabled(Rule::MissingReturnTypePrivateFunction) {
let return_type =
auto_return_type(function, checker.settings.target_version)
.map(|return_type| checker.generator().expr(&return_type));
let return_type = auto_return_type(function)
.and_then(|return_type| {
return_type.into_expression(
checker.importer(),
function.parameters.start(),
checker.semantic(),
checker.settings.target_version,
)
})
.map(|(return_type, edits)| {
(checker.generator().expr(&return_type), edits)
});
let mut diagnostic = Diagnostic::new(
MissingReturnTypePrivateFunction {
name: name.to_string(),
annotation: return_type.clone(),
annotation: return_type
.clone()
.map(|(return_type, ..)| return_type),
},
function.identifier(),
);
if let Some(return_type) = return_type {
diagnostic.set_fix(Fix::unsafe_edit(Edit::insertion(
format!(" -> {return_type}"),
function.parameters.range().end(),
)));
if let Some((return_type, edits)) = return_type {
diagnostic.set_fix(Fix::unsafe_edits(
Edit::insertion(
format!(" -> {return_type}"),
function.parameters.end(),
),
edits,
));
}
diagnostics.push(diagnostic);
}
Expand Down

0 comments on commit 6435e4e

Please sign in to comment.