diff --git a/crates/ruff_python_formatter/src/string/mod.rs b/crates/ruff_python_formatter/src/string/mod.rs index c41716427b34d..70b2288575ecd 100644 --- a/crates/ruff_python_formatter/src/string/mod.rs +++ b/crates/ruff_python_formatter/src/string/mod.rs @@ -1,13 +1,14 @@ use std::borrow::Cow; +use std::iter::FusedIterator; use bitflags::bitflags; use memchr::memchr2; use ruff_formatter::{format_args, write}; -use ruff_python_ast::AnyNodeRef; use ruff_python_ast::{ self as ast, Expr, ExprBytesLiteral, ExprFString, ExprStringLiteral, ExpressionRef, }; +use ruff_python_ast::{AnyNodeRef, StringLiteral}; use ruff_source_file::Locator; use ruff_text_size::{Ranged, TextLen, TextRange, TextSize}; @@ -68,31 +69,15 @@ impl<'a> AnyString<'a> { } /// Returns a vector of all the [`AnyStringPart`] of this string. - fn parts(self, quoting: Quoting) -> Vec> { + fn parts(self, quoting: Quoting) -> AnyStringPartsIter<'a> { match self { - Self::String(ExprStringLiteral { value, .. }) => value - .iter() - .map(|part| AnyStringPart::String { - part, - layout: StringLiteralKind::String, - }) - .collect(), - Self::Bytes(ExprBytesLiteral { value, .. }) => { - value.iter().map(AnyStringPart::Bytes).collect() + Self::String(ExprStringLiteral { value, .. }) => { + AnyStringPartsIter::String(value.iter()) + } + Self::Bytes(ExprBytesLiteral { value, .. }) => AnyStringPartsIter::Bytes(value.iter()), + Self::FString(ExprFString { value, .. }) => { + AnyStringPartsIter::FString(value.iter(), quoting) } - Self::FString(ExprFString { value, .. }) => value - .iter() - .map(|f_string_part| match f_string_part { - ast::FStringPart::Literal(string_literal) => AnyStringPart::String { - part: string_literal, - layout: StringLiteralKind::InImplicitlyConcatenatedFString(quoting), - }, - ast::FStringPart::FString(f_string) => AnyStringPart::FString { - part: f_string, - quoting, - }, - }) - .collect(), } } @@ -151,6 +136,46 @@ impl<'a> From<&AnyString<'a>> for ExpressionRef<'a> { } } +enum AnyStringPartsIter<'a> { + String(std::slice::Iter<'a, StringLiteral>), + Bytes(std::slice::Iter<'a, ast::BytesLiteral>), + FString(std::slice::Iter<'a, ast::FStringPart>, Quoting), +} + +impl<'a> Iterator for AnyStringPartsIter<'a> { + type Item = AnyStringPart<'a>; + + fn next(&mut self) -> Option { + let part = match self { + Self::String(inner) => { + let part = inner.next()?; + AnyStringPart::String { + part, + layout: StringLiteralKind::String, + } + } + Self::Bytes(inner) => AnyStringPart::Bytes(inner.next()?), + Self::FString(inner, quoting) => { + let part = inner.next()?; + match part { + ast::FStringPart::Literal(string_literal) => AnyStringPart::String { + part: string_literal, + layout: StringLiteralKind::InImplicitlyConcatenatedFString(*quoting), + }, + ast::FStringPart::FString(f_string) => AnyStringPart::FString { + part: f_string, + quoting: *quoting, + }, + } + } + }; + + Some(part) + } +} + +impl FusedIterator for AnyStringPartsIter<'_> {} + /// Represents any kind of string which is part of an implicitly concatenated /// string. This could be either a string, bytes or f-string. ///