Add autotyping-like return type inference for annotation rules (#8643)

## 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 https://github.com/astral-sh/ruff/issues/1640 (though we could
open a separate issue for referring parameter types).

Closes https://github.com/astral-sh/ruff/issues/8213.

## Test Plan

`cargo test`
This commit is contained in:
Charlie Marsh 2023-11-13 20:34:15 -08:00 committed by GitHub
parent 23c819b4b3
commit bf2cc3f520
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
18 changed files with 580 additions and 140 deletions

View file

@ -24,22 +24,41 @@ impl ResolvedPythonType {
(Self::TypeError, _) | (_, Self::TypeError) => Self::TypeError,
(Self::Unknown, _) | (_, Self::Unknown) => Self::Unknown,
(Self::Atom(a), Self::Atom(b)) => {
if a == b {
if a.is_subtype_of(b) {
Self::Atom(b)
} else if b.is_subtype_of(a) {
Self::Atom(a)
} else {
Self::Union(FxHashSet::from_iter([a, b]))
}
}
(Self::Atom(a), Self::Union(mut b)) => {
b.insert(a);
// If `a` is a subtype of any of the types in `b`, then `a` is
// redundant.
if !b.iter().any(|b_element| a.is_subtype_of(*b_element)) {
b.insert(a);
}
Self::Union(b)
}
(Self::Union(mut a), Self::Atom(b)) => {
a.insert(b);
// If `b` is a subtype of any of the types in `a`, then `b` is
// redundant.
if !a.iter().any(|a_element| b.is_subtype_of(*a_element)) {
a.insert(b);
}
Self::Union(a)
}
(Self::Union(mut a), Self::Union(b)) => {
a.extend(b);
for b_element in b {
// If `b_element` is a subtype of any of the types in `a`, then
// `b_element` is redundant.
if !a
.iter()
.any(|a_element| b_element.is_subtype_of(*a_element))
{
a.insert(b_element);
}
}
Self::Union(a)
}
}
@ -321,7 +340,7 @@ impl From<&Expr> for ResolvedPythonType {
/// such as strings, integers, floats, and containers. It cannot infer the
/// types of variables or expressions that are not statically known from
/// individual AST nodes alone.
#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum PythonType {
/// A string literal, such as `"hello"`.
String,
@ -345,8 +364,48 @@ pub enum PythonType {
Generator,
}
impl PythonType {
/// Returns `true` if `self` is a subtype of `other`.
fn is_subtype_of(self, other: Self) -> bool {
match (self, other) {
(PythonType::String, PythonType::String) => true,
(PythonType::Bytes, PythonType::Bytes) => true,
(PythonType::None, PythonType::None) => true,
(PythonType::Ellipsis, PythonType::Ellipsis) => true,
// The Numeric Tower (https://peps.python.org/pep-3141/)
(PythonType::Number(NumberLike::Bool), PythonType::Number(NumberLike::Bool)) => true,
(PythonType::Number(NumberLike::Integer), PythonType::Number(NumberLike::Integer)) => {
true
}
(PythonType::Number(NumberLike::Float), PythonType::Number(NumberLike::Float)) => true,
(PythonType::Number(NumberLike::Complex), PythonType::Number(NumberLike::Complex)) => {
true
}
(PythonType::Number(NumberLike::Bool), PythonType::Number(NumberLike::Integer)) => true,
(PythonType::Number(NumberLike::Bool), PythonType::Number(NumberLike::Float)) => true,
(PythonType::Number(NumberLike::Bool), PythonType::Number(NumberLike::Complex)) => true,
(PythonType::Number(NumberLike::Integer), PythonType::Number(NumberLike::Float)) => {
true
}
(PythonType::Number(NumberLike::Integer), PythonType::Number(NumberLike::Complex)) => {
true
}
(PythonType::Number(NumberLike::Float), PythonType::Number(NumberLike::Complex)) => {
true
}
// This simple type hierarchy doesn't support generics.
(PythonType::Dict, PythonType::Dict) => true,
(PythonType::List, PythonType::List) => true,
(PythonType::Set, PythonType::Set) => true,
(PythonType::Tuple, PythonType::Tuple) => true,
(PythonType::Generator, PythonType::Generator) => true,
_ => false,
}
}
}
/// A numeric type, or a type that can be trivially coerced to a numeric type.
#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum NumberLike {
/// An integer literal, such as `1` or `0x1`.
Integer,
@ -372,8 +431,6 @@ impl NumberLike {
#[cfg(test)]
mod tests {
use rustc_hash::FxHashSet;
use ruff_python_ast::Expr;
use ruff_python_parser::parse_expression;
@ -410,10 +467,7 @@ mod tests {
);
assert_eq!(
ResolvedPythonType::from(&parse("1 and True")),
ResolvedPythonType::Union(FxHashSet::from_iter([
PythonType::Number(NumberLike::Integer),
PythonType::Number(NumberLike::Bool)
]))
ResolvedPythonType::Atom(PythonType::Number(NumberLike::Integer))
);
// Binary operators.
@ -475,17 +529,11 @@ mod tests {
);
assert_eq!(
ResolvedPythonType::from(&parse("1 if True else 2.0")),
ResolvedPythonType::Union(FxHashSet::from_iter([
PythonType::Number(NumberLike::Integer),
PythonType::Number(NumberLike::Float)
]))
ResolvedPythonType::Atom(PythonType::Number(NumberLike::Float))
);
assert_eq!(
ResolvedPythonType::from(&parse("1 if True else False")),
ResolvedPythonType::Union(FxHashSet::from_iter([
PythonType::Number(NumberLike::Integer),
PythonType::Number(NumberLike::Bool)
]))
ResolvedPythonType::Atom(PythonType::Number(NumberLike::Integer))
);
}
}