From a01c378f764cc9a72b3a60ad0029610f0941e772 Mon Sep 17 00:00:00 2001 From: Dhruv Manilawala Date: Thu, 13 Jul 2023 11:01:18 +0530 Subject: [PATCH] Avoid allocation while resolving slice value Co-authored-by: Charlie Marsh --- crates/ruff/src/rules/ruff/typing.rs | 55 +++++++++++++++++----------- 1 file changed, 34 insertions(+), 21 deletions(-) diff --git a/crates/ruff/src/rules/ruff/typing.rs b/crates/ruff/src/rules/ruff/typing.rs index f96ea22d8ef22..35baa3f9ce881 100644 --- a/crates/ruff/src/rules/ruff/typing.rs +++ b/crates/ruff/src/rules/ruff/typing.rs @@ -1,3 +1,4 @@ +use itertools::Either::{Left, Right}; use rustpython_parser::ast::{self, Constant, Expr, Operator}; use ruff_python_ast::call_path::CallPath; @@ -52,12 +53,12 @@ fn is_known_type(call_path: &CallPath, minor_version: u32) -> bool { } } -/// Returns a vector of all the expressions in a slice. If the slice is not a -/// tuple, it will be wrapped in a vector. -fn resolve_slice_value(slice: &Expr) -> Vec<&Expr> { +/// Returns an iterator over the expressions in a slice. If the slice is not a +/// tuple, the iterator will only yield the slice. +fn resolve_slice_value(slice: &Expr) -> impl Iterator { match slice { - Expr::Tuple(ast::ExprTuple { elts: elements, .. }) => elements.iter().collect(), - _ => vec![slice], + Expr::Tuple(ast::ExprTuple { elts: elements, .. }) => Left(elements.iter()), + _ => Right(std::iter::once(slice)), } } @@ -75,12 +76,14 @@ enum TypingTarget<'a> { /// Forward reference to a type e.g., `"List[str]"`. ForwardReference(Expr), - /// A `typing.Union` type or `|` separated types e.g., `Union[int, str]` - /// or `int | str`. - Union(Vec<&'a Expr>), + /// A `typing.Union` type e.g., `Union[int, str]`. + Union(&'a Expr), + + /// A PEP 604 union type e.g., `int | str`. + PEP604Union(&'a Expr), /// A `typing.Literal` type e.g., `Literal[1, 2, 3]`. - Literal(Vec<&'a Expr>), + Literal(&'a Expr), /// A `typing.Optional` type e.g., `Optional[int]`. Optional(&'a Expr), @@ -110,13 +113,13 @@ impl<'a> TypingTarget<'a> { if semantic.match_typing_expr(value, "Optional") { Some(TypingTarget::Optional(slice.as_ref())) } else if semantic.match_typing_expr(value, "Literal") { - Some(TypingTarget::Literal(resolve_slice_value(slice))) + Some(TypingTarget::Literal(slice)) } else if semantic.match_typing_expr(value, "Union") { - Some(TypingTarget::Union(resolve_slice_value(slice))) + Some(TypingTarget::Union(slice)) } else if semantic.match_typing_expr(value, "Annotated") { resolve_slice_value(slice.as_ref()) - .first() - .map(|&expr| TypingTarget::Annotated(expr)) + .next() + .map(TypingTarget::Annotated) } else { semantic.resolve_call_path(value).map_or( // If we can't resolve the call path, it must be defined @@ -132,9 +135,7 @@ impl<'a> TypingTarget<'a> { ) } } - Expr::BinOp(..) => Some(TypingTarget::Union( - PEP604UnionIterator::new(expr).collect(), - )), + Expr::BinOp(..) => Some(TypingTarget::PEP604Union(expr)), Expr::Constant(ast::ExprConstant { value: Constant::None, .. @@ -179,7 +180,7 @@ impl<'a> TypingTarget<'a> { | TypingTarget::Object | TypingTarget::Unknown => true, TypingTarget::Known => false, - TypingTarget::Literal(elements) => elements.iter().any(|element| { + TypingTarget::Literal(slice) => resolve_slice_value(slice).any(|element| { // Literal can only contain `None`, a literal value, other `Literal` // or an enum value. match TypingTarget::try_from_expr(element, semantic, locator, minor_version) { @@ -190,17 +191,23 @@ impl<'a> TypingTarget<'a> { _ => false, } }), - TypingTarget::Union(elements) => elements.iter().any(|element| { + TypingTarget::Union(slice) => resolve_slice_value(slice).any(|element| { TypingTarget::try_from_expr(element, semantic, locator, minor_version) .map_or(true, |new_target| { new_target.contains_none(semantic, locator, minor_version) }) }), - TypingTarget::Annotated(element) => { + TypingTarget::PEP604Union(expr) => PEP604UnionIterator::new(expr).any(|element| { TypingTarget::try_from_expr(element, semantic, locator, minor_version) .map_or(true, |new_target| { new_target.contains_none(semantic, locator, minor_version) }) + }), + TypingTarget::Annotated(expr) => { + TypingTarget::try_from_expr(expr, semantic, locator, minor_version) + .map_or(true, |new_target| { + new_target.contains_none(semantic, locator, minor_version) + }) } TypingTarget::ForwardReference(expr) => { TypingTarget::try_from_expr(expr, semantic, locator, minor_version) @@ -226,17 +233,23 @@ impl<'a> TypingTarget<'a> { | TypingTarget::Object | TypingTarget::Known | TypingTarget::Unknown => false, - TypingTarget::Union(elements) => elements.iter().any(|element| { + TypingTarget::Union(slice) => resolve_slice_value(slice).any(|element| { TypingTarget::try_from_expr(element, semantic, locator, minor_version) .map_or(true, |new_target| { new_target.contains_any(semantic, locator, minor_version) }) }), - TypingTarget::Annotated(element) | TypingTarget::Optional(element) => { + TypingTarget::PEP604Union(expr) => PEP604UnionIterator::new(expr).any(|element| { TypingTarget::try_from_expr(element, semantic, locator, minor_version) .map_or(true, |new_target| { new_target.contains_any(semantic, locator, minor_version) }) + }), + TypingTarget::Annotated(expr) | TypingTarget::Optional(expr) => { + TypingTarget::try_from_expr(expr, semantic, locator, minor_version) + .map_or(true, |new_target| { + new_target.contains_any(semantic, locator, minor_version) + }) } TypingTarget::ForwardReference(expr) => { TypingTarget::try_from_expr(expr, semantic, locator, minor_version)