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:


128414cd95/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.
This commit is contained in:
Dhruv Manilawala 2024-05-13 21:22:03 +05:30 committed by GitHub
parent 10b85a0f07
commit 0dc130e841
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 129 additions and 81 deletions

View file

@ -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<Self::Item> {
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<usize>) {
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<'_> {}