Skip to content

Commit

Permalink
Remove an unwrap from unnecessary_literal_union.rs (#9404)
Browse files Browse the repository at this point in the history
  • Loading branch information
@charliermarsh
charliermarsh committed Jan 5, 2024
1 parent 59078c5 commit 62eca33
Show file tree
Hide file tree
Showing 2 changed files with 78 additions and 123 deletions.
159 changes: 66 additions & 93 deletions crates/ruff_linter/src/rules/flake8_pyi/rules/unnecessary_literal_union.rs
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
use ast::Operator;
use ruff_diagnostics::{Diagnostic, Edit, Fix, FixAvailability, Violation};
use ruff_macros::{derive_message_formats, violation};
use ruff_python_ast::{self as ast, Expr};
use ruff_python_ast::helpers::pep_604_union;
use ruff_python_ast::{self as ast, Expr, ExprContext};
use ruff_python_semantic::analyze::typing::traverse_union;
use ruff_text_size::{Ranged, TextRange};

Expand Down Expand Up @@ -48,61 +48,12 @@ impl Violation for UnnecessaryLiteralUnion {
}
}

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(),
})
})
}

fn make_union(subscript: &ast::ExprSubscript, exprs: Vec<&Expr>) -> Expr {
Expr::Subscript(ast::ExprSubscript {
value: subscript.value.clone(),
slice: Box::new(Expr::Tuple(ast::ExprTuple {
elts: exprs.into_iter().map(|expr| (*expr).clone()).collect(),
range: TextRange::default(),
ctx: ast::ExprContext::Load,
})),
range: TextRange::default(),
ctx: ast::ExprContext::Load,
})
}

fn make_literal_expr(subscript: Option<Expr>, exprs: Vec<&Expr>) -> Expr {
let use_subscript = if let subscript @ Some(_) = subscript {
subscript.unwrap().clone()
} else {
Expr::Name(ast::ExprName {
id: "Literal".to_string(),
range: TextRange::default(),
ctx: ast::ExprContext::Load,
})
};
Expr::Subscript(ast::ExprSubscript {
value: Box::new(use_subscript),
slice: Box::new(Expr::Tuple(ast::ExprTuple {
elts: exprs.into_iter().map(|expr| (*expr).clone()).collect(),
range: TextRange::default(),
ctx: ast::ExprContext::Load,
})),
range: TextRange::default(),
ctx: ast::ExprContext::Load,
})
}

/// PYI030
pub(crate) fn unnecessary_literal_union<'a>(checker: &mut Checker, expr: &'a Expr) {
let mut literal_exprs = Vec::new();
let mut other_exprs = Vec::new();

// for the sake of consistency and correctness, we'll use the first `Literal` subscript attribute
// to construct the fix
// For the sake of consistency, use the first `Literal` subscript to construct the fix.
let mut literal_subscript = None;
let mut total_literals = 0;

Expand All @@ -113,7 +64,7 @@ pub(crate) fn unnecessary_literal_union<'a>(checker: &mut Checker, expr: &'a Exp
total_literals += 1;

if literal_subscript.is_none() {
literal_subscript = Some(*value.clone());
literal_subscript = Some(value.as_ref());
}

// flatten already-unioned literals to later union again
Expand Down Expand Up @@ -147,48 +98,70 @@ pub(crate) fn unnecessary_literal_union<'a>(checker: &mut Checker, expr: &'a Exp
return;
}

// Raise a violation if more than one.
if total_literals > 1 {
let literal_members: Vec<String> = literal_exprs
.clone()
.into_iter()
.map(|expr| checker.locator().slice(expr).to_string())
.collect();

let mut diagnostic = Diagnostic::new(
UnnecessaryLiteralUnion {
members: literal_members.clone(),
},
expr.range(),
);

if checker.settings.preview.is_enabled() {
let literals =
make_literal_expr(literal_subscript, literal_exprs.into_iter().collect());

if other_exprs.is_empty() {
// if the union is only literals, we just replace the whole thing with a single literal
diagnostic.set_fix(Fix::safe_edit(Edit::range_replacement(
checker.generator().expr(&literals),
expr.range(),
)));
} else {
let mut expr_vec: Vec<&Expr> = other_exprs.clone().into_iter().collect();
expr_vec.insert(0, &literals);
// If there are no literal members, we don't need to do anything.
let Some(literal_subscript) = literal_subscript else {
return;
};
if total_literals == 0 || total_literals == 1 {
return;
}

let content = if let Some(subscript) = union_subscript {
checker.generator().expr(&make_union(subscript, expr_vec))
} else {
checker.generator().expr(&concatenate_bin_ors(expr_vec))
};
let mut diagnostic = Diagnostic::new(
UnnecessaryLiteralUnion {
members: literal_exprs
.iter()
.map(|expr| checker.locator().slice(expr).to_string())
.collect(),
},
expr.range(),
);

if checker.settings.preview.is_enabled() {
let literal = Expr::Subscript(ast::ExprSubscript {
value: Box::new(literal_subscript.clone()),
slice: Box::new(Expr::Tuple(ast::ExprTuple {
elts: literal_exprs.into_iter().cloned().collect(),
range: TextRange::default(),
ctx: ExprContext::Load,
})),
range: TextRange::default(),
ctx: ExprContext::Load,
});

if other_exprs.is_empty() {
// if the union is only literals, we just replace the whole thing with a single literal
diagnostic.set_fix(Fix::safe_edit(Edit::range_replacement(
checker.generator().expr(&literal),
expr.range(),
)));
} else {
let elts: Vec<Expr> = std::iter::once(literal)
.chain(other_exprs.into_iter().cloned())
.collect();

let content = if let Some(union) = union_subscript {
checker
.generator()
.expr(&Expr::Subscript(ast::ExprSubscript {
value: union.value.clone(),
slice: Box::new(Expr::Tuple(ast::ExprTuple {
elts,
range: TextRange::default(),
ctx: ExprContext::Load,
})),
range: TextRange::default(),
ctx: ExprContext::Load,
}))
} else {
checker.generator().expr(&pep_604_union(&elts))
};

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

checker.diagnostics.push(diagnostic);
}

checker.diagnostics.push(diagnostic);
}
42 changes: 12 additions & 30 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,6 +1,7 @@
use ast::{ExprContext, Operator};
use ast::ExprContext;
use ruff_diagnostics::{Diagnostic, Edit, Fix, FixAvailability, Violation};
use ruff_macros::{derive_message_formats, violation};
use ruff_python_ast::helpers::pep_604_union;
use ruff_python_ast::{self as ast, Expr};
use ruff_python_semantic::analyze::typing::traverse_union;
use ruff_text_size::{Ranged, TextRange};
Expand Down Expand Up @@ -50,19 +51,6 @@ impl Violation for UnnecessaryTypeUnion {
}
}

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
pub(crate) fn unnecessary_type_union<'a>(checker: &mut Checker, union: &'a Expr) {
// The `|` operator isn't always safe to allow to runtime-evaluated annotations.
Expand Down Expand Up @@ -95,7 +83,7 @@ pub(crate) fn unnecessary_type_union<'a>(checker: &mut Checker, union: &'a Expr)
.resolve_call_path(unwrapped.value.as_ref())
.is_some_and(|call_path| matches!(call_path.as_slice(), ["" | "builtins", "type"]))
{
type_exprs.push(&unwrapped.slice);
type_exprs.push(unwrapped.slice.as_ref());
} else {
other_exprs.push(expr);
}
Expand All @@ -108,7 +96,7 @@ pub(crate) fn unnecessary_type_union<'a>(checker: &mut Checker, union: &'a Expr)
let type_members: Vec<String> = type_exprs
.clone()
.into_iter()
.map(|type_expr| checker.locator().slice(type_expr.as_ref()).to_string())
.map(|type_expr| checker.locator().slice(type_expr).to_string())
.collect();

let mut diagnostic = Diagnostic::new(
Expand Down Expand Up @@ -171,31 +159,25 @@ pub(crate) fn unnecessary_type_union<'a>(checker: &mut Checker, union: &'a Expr)
checker.generator().expr(&union)
}
} else {
let types = &Expr::Subscript(ast::ExprSubscript {
let elts: Vec<Expr> = type_exprs.into_iter().cloned().collect();
let types = 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(),
)),
slice: Box::new(pep_604_union(&elts)),
ctx: ExprContext::Load,
range: TextRange::default(),
});

if other_exprs.is_empty() {
checker.generator().expr(types)
checker.generator().expr(&types)
} else {
let mut exprs = Vec::new();
exprs.push(types);
exprs.extend(other_exprs);

checker.generator().expr(&concatenate_bin_ors(exprs))
let elts: Vec<Expr> = std::iter::once(types)
.chain(other_exprs.into_iter().cloned())
.collect();
checker.generator().expr(&pep_604_union(&elts))
}
};

Expand Down

0 comments on commit 62eca33

Please sign in to comment.