Skip to content

Commit

Permalink
Add autofix for PYI055 (#7886)
Browse files Browse the repository at this point in the history
  • Loading branch information
@diceroll123
diceroll123 authored Oct 13, 2023
1 parent f08a5f6 commit 1e184e6
Show file tree
Hide file tree
Showing 5 changed files with 277 additions and 32 deletions.
10 changes: 9 additions & 1 deletion crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI055.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,4 +28,12 @@ def func(arg: type[int, float] | str) -> None:

def func():
# PYI055
item: type[requests_mock.Mocker] | type[httpretty] = requests_mock.Mocker
x: type[requests_mock.Mocker] | type[httpretty] | type[str] = requests_mock.Mocker
y: Union[type[requests_mock.Mocker], type[httpretty], type[str]] = requests_mock.Mocker


def func():
from typing import Union as U

# PYI055
x: Union[type[requests_mock.Mocker], type[httpretty], type[str]] = requests_mock.Mocker
3 changes: 2 additions & 1 deletion crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI055.pyi
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,4 +21,5 @@ item: type[requests_mock.Mocker] | type[httpretty] = requests_mock.Mocker

def func():
# PYI055
item: type[requests_mock.Mocker] | type[httpretty] = requests_mock.Mocker
item: type[requests_mock.Mocker] | type[httpretty] | type[str] = requests_mock.Mocker
item2: Union[type[requests_mock.Mocker], type[httpretty], type[str]] = requests_mock.Mocker
121 changes: 104 additions & 17 deletions crates/ruff_linter/src/rules/flake8_pyi/rules/unnecessary_type_union.rs
View file
Original file line number Diff line number Diff line change
@@ -1,16 +1,17 @@
use ruff_diagnostics::{Diagnostic, Violation};
use ast::{ExprContext, Operator};
use ruff_diagnostics::{Diagnostic, Edit, Fix, FixAvailability, Violation};
use ruff_macros::{derive_message_formats, violation};
use ruff_python_ast::Expr;
use ruff_text_size::Ranged;
use ruff_python_ast::{self as ast, Expr};
use ruff_text_size::{Ranged, TextRange};

use crate::{checkers::ast::Checker, rules::flake8_pyi::helpers::traverse_union};

/// ## What it does
/// Checks for the presence of multiple `type`s in a union.
///
/// ## Why is this bad?
/// The `type` built-in function accepts unions, and it is
/// clearer to explicitly specify them as a single `type`.
/// The `type` built-in function accepts unions, and it is clearer to
/// explicitly specify them as a single `type`.
///
/// ## Example
/// ```python
Expand All @@ -28,6 +29,8 @@ pub struct UnnecessaryTypeUnion {
}

impl Violation for UnnecessaryTypeUnion {
const FIX_AVAILABILITY: FixAvailability = FixAvailability::Sometimes;

#[derive_message_formats]
fn message(&self) -> String {
let union_str = if self.is_pep604_union {
Expand All @@ -40,6 +43,23 @@ impl Violation for UnnecessaryTypeUnion {
"Multiple `type` members in a union. Combine them into one, e.g., `type[{union_str}]`."
)
}

fn fix_title(&self) -> Option<String> {
Some(format!("Combine multiple `type` members"))
}
}

fn concatenate_bin_ors(exprs: Vec<&Expr>) -> Expr {
let mut exprs = exprs.into_iter();
let first = exprs.next().unwrap();
exprs.fold((*first).clone(), |acc, expr| {
Expr::BinOp(ast::ExprBinOp {
left: Box::new(acc),
op: Operator::BitOr,
right: Box::new((*expr).clone()),
range: TextRange::default(),
})
})
}

/// PYI055
Expand All @@ -49,14 +69,17 @@ pub(crate) fn unnecessary_type_union<'a>(checker: &mut Checker, union: &'a Expr)
return;
}

let mut type_exprs = Vec::new();

// Check if `union` is a PEP604 union (e.g. `float | int`) or a `typing.Union[float, int]`
let is_pep604_union = !union.as_subscript_expr().is_some_and(|subscript| {
checker
let subscript = union.as_subscript_expr();
if subscript.is_some_and(|subscript| {
!checker
.semantic()
.match_typing_expr(&subscript.value, "Union")
});
}) {
return;
}

let mut type_exprs = Vec::new();

let mut collect_type_exprs = |expr: &'a Expr, _| {
let Some(subscript) = expr.as_subscript_expr() else {
Expand All @@ -74,15 +97,79 @@ pub(crate) fn unnecessary_type_union<'a>(checker: &mut Checker, union: &'a Expr)
traverse_union(&mut collect_type_exprs, checker.semantic(), union, None);

if type_exprs.len() > 1 {
checker.diagnostics.push(Diagnostic::new(
let type_members: Vec<String> = type_exprs
.clone()
.into_iter()
.map(|type_expr| checker.locator().slice(type_expr.as_ref()).to_string())
.collect();

let mut diagnostic = Diagnostic::new(
UnnecessaryTypeUnion {
members: type_exprs
.into_iter()
.map(|type_expr| checker.locator().slice(type_expr.as_ref()).to_string())
.collect(),
is_pep604_union,
members: type_members.clone(),
is_pep604_union: subscript.is_none(),
},
union.range(),
));
);

if checker.semantic().is_builtin("type") {
let content = if let Some(subscript) = subscript {
checker
.generator()
.expr(&Expr::Subscript(ast::ExprSubscript {
value: Box::new(Expr::Name(ast::ExprName {
id: "type".into(),
ctx: ExprContext::Load,
range: TextRange::default(),
})),
slice: Box::new(Expr::Subscript(ast::ExprSubscript {
value: subscript.value.clone(),
slice: Box::new(Expr::Tuple(ast::ExprTuple {
elts: type_members
.into_iter()
.map(|type_member| {
Expr::Name(ast::ExprName {
id: type_member,
ctx: ExprContext::Load,
range: TextRange::default(),
})
})
.collect(),
ctx: ExprContext::Load,
range: TextRange::default(),
})),
ctx: ExprContext::Load,
range: TextRange::default(),
})),
ctx: ExprContext::Load,
range: TextRange::default(),
}))
} else {
checker
.generator()
.expr(&Expr::Subscript(ast::ExprSubscript {
value: Box::new(Expr::Name(ast::ExprName {
id: "type".into(),
ctx: ExprContext::Load,
range: TextRange::default(),
})),
slice: Box::new(concatenate_bin_ors(
type_exprs
.clone()
.into_iter()
.map(std::convert::AsRef::as_ref)
.collect(),
)),
ctx: ExprContext::Load,
range: TextRange::default(),
}))
};

diagnostic.set_fix(Fix::safe_edit(Edit::range_replacement(
content,
union.range(),
)));
}

checker.diagnostics.push(diagnostic);
}
}
52 changes: 49 additions & 3 deletions ...c/rules/flake8_pyi/snapshots/ruff_linter__rules__flake8_pyi__tests__PYI055_PYI055.py.snap
View file
Original file line number Diff line number Diff line change
@@ -1,12 +1,58 @@
---
source: crates/ruff_linter/src/rules/flake8_pyi/mod.rs
---
PYI055.py:31:11: PYI055 Multiple `type` members in a union. Combine them into one, e.g., `type[requests_mock.Mocker | httpretty]`.
PYI055.py:31:8: PYI055 [*] Multiple `type` members in a union. Combine them into one, e.g., `type[requests_mock.Mocker | httpretty | str]`.
|
29 | def func():
30 | # PYI055
31 | item: type[requests_mock.Mocker] | type[httpretty] = requests_mock.Mocker
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ PYI055
31 | x: type[requests_mock.Mocker] | type[httpretty] | type[str] = requests_mock.Mocker
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ PYI055
32 | y: Union[type[requests_mock.Mocker], type[httpretty], type[str]] = requests_mock.Mocker
|
= help: Combine multiple `type` members

ℹ Fix
28 28 |
29 29 | def func():
30 30 | # PYI055
31 |- x: type[requests_mock.Mocker] | type[httpretty] | type[str] = requests_mock.Mocker
31 |+ x: type[requests_mock.Mocker | httpretty | str] = requests_mock.Mocker
32 32 | y: Union[type[requests_mock.Mocker], type[httpretty], type[str]] = requests_mock.Mocker
33 33 |
34 34 |

PYI055.py:32:8: PYI055 [*] Multiple `type` members in a union. Combine them into one, e.g., `type[Union[requests_mock.Mocker, httpretty, str]]`.
|
30 | # PYI055
31 | x: type[requests_mock.Mocker] | type[httpretty] | type[str] = requests_mock.Mocker
32 | y: Union[type[requests_mock.Mocker], type[httpretty], type[str]] = requests_mock.Mocker
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ PYI055
|
= help: Combine multiple `type` members

ℹ Fix
29 29 | def func():
30 30 | # PYI055
31 31 | x: type[requests_mock.Mocker] | type[httpretty] | type[str] = requests_mock.Mocker
32 |- y: Union[type[requests_mock.Mocker], type[httpretty], type[str]] = requests_mock.Mocker
32 |+ y: type[Union[requests_mock.Mocker, httpretty, str]] = requests_mock.Mocker
33 33 |
34 34 |
35 35 | def func():

PYI055.py:39:8: PYI055 [*] Multiple `type` members in a union. Combine them into one, e.g., `type[Union[requests_mock.Mocker, httpretty, str]]`.
|
38 | # PYI055
39 | x: Union[type[requests_mock.Mocker], type[httpretty], type[str]] = requests_mock.Mocker
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ PYI055
|
= help: Combine multiple `type` members

ℹ Fix
36 36 | from typing import Union as U
37 37 |
38 38 | # PYI055
39 |- x: Union[type[requests_mock.Mocker], type[httpretty], type[str]] = requests_mock.Mocker
39 |+ x: type[Union[requests_mock.Mocker, httpretty, str]] = requests_mock.Mocker


Loading

0 comments on commit 1e184e6

Please sign in to comment.