From 0dc130e841c6c7cd3df429ce064bae1ff2954d54 Mon Sep 17 00:00:00 2001 From: Dhruv Manilawala Date: Mon, 13 May 2024 21:22:03 +0530 Subject: [PATCH] Add `Iterator` impl for `StringLike` parts (#11399) ## Summary This PR adds support to iterate over each part of a string-like expression. This similar to the one in the formatter: https://github.com/astral-sh/ruff/blob/128414cd9507adc45584189f025363793dd7cd94/crates/ruff_python_formatter/src/string/any.rs#L121-L125 Although I don't think it's a 1-1 replacement in the formatter because the one implemented in the formatter has another information for certain variants (as can be seen for `FString`). The main motivation for this is to avoid duplication for rules which work only on the parts of the string and doesn't require any information from the parent node. Here, the parent node being the expression node which could be an implicitly concatenated string. This PR also updates certain rule implementation to make use of this and avoids logic duplication. --- .../rules/avoidable_escaped_quote.rs | 23 +---- .../rules/check_string_quotes.rs | 6 +- .../rules/unnecessary_escaped_quote.rs | 37 +++---- .../ruff/rules/ambiguous_unicode_character.rs | 48 ++++------ crates/ruff_python_ast/src/expression.rs | 96 ++++++++++++++++++- 5 files changed, 129 insertions(+), 81 deletions(-) diff --git a/crates/ruff_linter/src/rules/flake8_quotes/rules/avoidable_escaped_quote.rs b/crates/ruff_linter/src/rules/flake8_quotes/rules/avoidable_escaped_quote.rs index d66a85ccb2..d17307239c 100644 --- a/crates/ruff_linter/src/rules/flake8_quotes/rules/avoidable_escaped_quote.rs +++ b/crates/ruff_linter/src/rules/flake8_quotes/rules/avoidable_escaped_quote.rs @@ -64,28 +64,15 @@ pub(crate) fn avoidable_escaped_quote(checker: &mut Checker, string_like: String let mut rule_checker = AvoidableEscapedQuoteChecker::new(checker.locator(), checker.settings); - match string_like { - StringLike::String(expr) => { - for string_literal in &expr.value { + for part in string_like.parts() { + match part { + ast::StringLikePart::String(string_literal) => { rule_checker.visit_string_literal(string_literal); } - } - StringLike::Bytes(expr) => { - for bytes_literal in &expr.value { + ast::StringLikePart::Bytes(bytes_literal) => { rule_checker.visit_bytes_literal(bytes_literal); } - } - StringLike::FString(expr) => { - for part in &expr.value { - match part { - ast::FStringPart::Literal(string_literal) => { - rule_checker.visit_string_literal(string_literal); - } - ast::FStringPart::FString(f_string) => { - rule_checker.visit_f_string(f_string); - } - } - } + ast::StringLikePart::FString(f_string) => rule_checker.visit_f_string(f_string), } } diff --git a/crates/ruff_linter/src/rules/flake8_quotes/rules/check_string_quotes.rs b/crates/ruff_linter/src/rules/flake8_quotes/rules/check_string_quotes.rs index 5ccffe6316..3a2db2504e 100644 --- a/crates/ruff_linter/src/rules/flake8_quotes/rules/check_string_quotes.rs +++ b/crates/ruff_linter/src/rules/flake8_quotes/rules/check_string_quotes.rs @@ -454,11 +454,7 @@ pub(crate) fn check_string_quotes(checker: &mut Checker, string_like: StringLike return; } - let ranges: Vec = match string_like { - StringLike::String(node) => node.value.iter().map(Ranged::range).collect(), - StringLike::Bytes(node) => node.value.iter().map(Ranged::range).collect(), - StringLike::FString(node) => node.value.iter().map(Ranged::range).collect(), - }; + let ranges: Vec<_> = string_like.parts().map(|part| part.range()).collect(); if checker.semantic().in_pep_257_docstring() { if checker.enabled(Rule::BadQuotesDocstring) { diff --git a/crates/ruff_linter/src/rules/flake8_quotes/rules/unnecessary_escaped_quote.rs b/crates/ruff_linter/src/rules/flake8_quotes/rules/unnecessary_escaped_quote.rs index f0da47b156..68e7ed67f9 100644 --- a/crates/ruff_linter/src/rules/flake8_quotes/rules/unnecessary_escaped_quote.rs +++ b/crates/ruff_linter/src/rules/flake8_quotes/rules/unnecessary_escaped_quote.rs @@ -53,39 +53,30 @@ pub(crate) fn unnecessary_escaped_quote(checker: &mut Checker, string_like: Stri let locator = checker.locator(); - match string_like { - StringLike::String(expr) => { - for string in &expr.value { + for part in string_like.parts() { + match part { + ast::StringLikePart::String(string_literal) => { if let Some(diagnostic) = check_string_or_bytes( locator, - string.range(), - AnyStringFlags::from(string.flags), + string_literal.range(), + AnyStringFlags::from(string_literal.flags), ) { checker.diagnostics.push(diagnostic); } } - } - StringLike::Bytes(expr) => { - for bytes in &expr.value { - if let Some(diagnostic) = - check_string_or_bytes(locator, bytes.range(), AnyStringFlags::from(bytes.flags)) - { + ast::StringLikePart::Bytes(bytes_literal) => { + if let Some(diagnostic) = check_string_or_bytes( + locator, + bytes_literal.range(), + AnyStringFlags::from(bytes_literal.flags), + ) { checker.diagnostics.push(diagnostic); } } - } - StringLike::FString(expr) => { - for part in &expr.value { - if let Some(diagnostic) = match part { - ast::FStringPart::Literal(string) => check_string_or_bytes( - locator, - string.range(), - AnyStringFlags::from(string.flags), - ), - ast::FStringPart::FString(f_string) => check_f_string(locator, f_string), - } { + ast::StringLikePart::FString(f_string) => { + if let Some(diagnostic) = check_f_string(locator, f_string) { checker.diagnostics.push(diagnostic); - }; + } } } } diff --git a/crates/ruff_linter/src/rules/ruff/rules/ambiguous_unicode_character.rs b/crates/ruff_linter/src/rules/ruff/rules/ambiguous_unicode_character.rs index 77c828848f..2dc499c9d3 100644 --- a/crates/ruff_linter/src/rules/ruff/rules/ambiguous_unicode_character.rs +++ b/crates/ruff_linter/src/rules/ruff/rules/ambiguous_unicode_character.rs @@ -192,48 +192,32 @@ pub(crate) fn ambiguous_unicode_character_string(checker: &mut Checker, string_l Context::String }; - match string_like { - StringLike::String(node) => { - for literal in &node.value { - let text = checker.locator().slice(literal); + for part in string_like.parts() { + match part { + ast::StringLikePart::String(string_literal) => { + let text = checker.locator().slice(string_literal); ambiguous_unicode_character( &mut checker.diagnostics, text, - literal.range(), + string_literal.range(), context, checker.settings, ); } - } - StringLike::FString(node) => { - for part in &node.value { - match part { - ast::FStringPart::Literal(literal) => { - let text = checker.locator().slice(literal); - ambiguous_unicode_character( - &mut checker.diagnostics, - text, - literal.range(), - context, - checker.settings, - ); - } - ast::FStringPart::FString(f_string) => { - for literal in f_string.literals() { - let text = checker.locator().slice(literal); - ambiguous_unicode_character( - &mut checker.diagnostics, - text, - literal.range(), - context, - checker.settings, - ); - } - } + ast::StringLikePart::Bytes(_) => {} + ast::StringLikePart::FString(f_string) => { + for literal in f_string.literals() { + let text = checker.locator().slice(literal); + ambiguous_unicode_character( + &mut checker.diagnostics, + text, + literal.range(), + context, + checker.settings, + ); } } } - StringLike::Bytes(_) => (), } } diff --git a/crates/ruff_python_ast/src/expression.rs b/crates/ruff_python_ast/src/expression.rs index 1ef7247f5a..3e19d5b3fd 100644 --- a/crates/ruff_python_ast/src/expression.rs +++ b/crates/ruff_python_ast/src/expression.rs @@ -1,3 +1,5 @@ +use std::iter::FusedIterator; + use ruff_text_size::{Ranged, TextRange}; use crate::AnyNodeRef; @@ -394,9 +396,8 @@ impl LiteralExpressionRef<'_> { } } -/// An enum that holds a reference to a string-like literal from the AST. -/// This includes string literals, bytes literals, and the literal parts of -/// f-strings. +/// An enum that holds a reference to a string-like expression from the AST. This includes string +/// literals, bytes literals, and f-strings. #[derive(Copy, Clone, Debug, PartialEq)] pub enum StringLike<'a> { String(&'a ast::ExprStringLiteral), @@ -404,6 +405,17 @@ pub enum StringLike<'a> { FString(&'a ast::ExprFString), } +impl<'a> StringLike<'a> { + /// Returns an iterator over the [`StringLikePart`] contained in this string-like expression. + pub fn parts(&self) -> StringLikePartIter<'_> { + match self { + StringLike::String(expr) => StringLikePartIter::String(expr.value.iter()), + StringLike::Bytes(expr) => StringLikePartIter::Bytes(expr.value.iter()), + StringLike::FString(expr) => StringLikePartIter::FString(expr.value.iter()), + } + } +} + impl<'a> From<&'a ast::ExprStringLiteral> for StringLike<'a> { fn from(value: &'a ast::ExprStringLiteral) -> Self { StringLike::String(value) @@ -431,3 +443,81 @@ impl Ranged for StringLike<'_> { } } } + +/// An enum that holds a reference to an individual part of a string-like expression. +#[derive(Copy, Clone, Debug, PartialEq)] +pub enum StringLikePart<'a> { + String(&'a ast::StringLiteral), + Bytes(&'a ast::BytesLiteral), + FString(&'a ast::FString), +} + +impl<'a> From<&'a ast::StringLiteral> for StringLikePart<'a> { + fn from(value: &'a ast::StringLiteral) -> Self { + StringLikePart::String(value) + } +} + +impl<'a> From<&'a ast::BytesLiteral> for StringLikePart<'a> { + fn from(value: &'a ast::BytesLiteral) -> Self { + StringLikePart::Bytes(value) + } +} + +impl<'a> From<&'a ast::FString> for StringLikePart<'a> { + fn from(value: &'a ast::FString) -> Self { + StringLikePart::FString(value) + } +} + +impl Ranged for StringLikePart<'_> { + fn range(&self) -> TextRange { + match self { + StringLikePart::String(part) => part.range(), + StringLikePart::Bytes(part) => part.range(), + StringLikePart::FString(part) => part.range(), + } + } +} + +/// An iterator over all the [`StringLikePart`] of a string-like expression. +/// +/// This is created by the [`StringLike::parts`] method. +pub enum StringLikePartIter<'a> { + String(std::slice::Iter<'a, ast::StringLiteral>), + Bytes(std::slice::Iter<'a, ast::BytesLiteral>), + FString(std::slice::Iter<'a, ast::FStringPart>), +} + +impl<'a> Iterator for StringLikePartIter<'a> { + type Item = StringLikePart<'a>; + + fn next(&mut self) -> Option { + let part = match self { + StringLikePartIter::String(inner) => StringLikePart::String(inner.next()?), + StringLikePartIter::Bytes(inner) => StringLikePart::Bytes(inner.next()?), + StringLikePartIter::FString(inner) => { + let part = inner.next()?; + match part { + ast::FStringPart::Literal(string_literal) => { + StringLikePart::String(string_literal) + } + ast::FStringPart::FString(f_string) => StringLikePart::FString(f_string), + } + } + }; + + Some(part) + } + + fn size_hint(&self) -> (usize, Option) { + match self { + StringLikePartIter::String(inner) => inner.size_hint(), + StringLikePartIter::Bytes(inner) => inner.size_hint(), + StringLikePartIter::FString(inner) => inner.size_hint(), + } + } +} + +impl FusedIterator for StringLikePartIter<'_> {} +impl ExactSizeIterator for StringLikePartIter<'_> {}