Skip to content

Commit

Permalink
Add as_slice method for all string nodes (#9111)
Browse files Browse the repository at this point in the history
This PR adds a `as_slice` method to all the string nodes which returns
all the parts of the nodes as a slice. This will be useful in the next
PR to split the string formatting to use this method to extract the
_single node_ or _implicitly concanated nodes_.
  • Loading branch information
dhruvmanila committed Dec 13, 2023
1 parent cb99815 commit 18452cf
Show file tree
Hide file tree
Showing 13 changed files with 120 additions and 64 deletions.
2 changes: 1 addition & 1 deletion crates/ruff_linter/src/checkers/ast/analyze/expression.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1277,7 +1277,7 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
}
Expr::StringLiteral(ast::ExprStringLiteral { value, .. }) => {
if checker.enabled(Rule::UnicodeKindPrefix) {
for string_part in value.parts() {
for string_part in value {
pyupgrade::rules::unicode_kind_prefix(checker, string_part);
}
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -57,7 +57,7 @@ impl Violation for HardcodedSQLExpression {
/// becomes `foobar {x}baz`.
fn concatenated_f_string(expr: &ast::ExprFString, locator: &Locator) -> String {
expr.value
.parts()
.iter()
.filter_map(|part| {
raw_contents(locator.slice(part)).map(|s| s.escape_default().to_string())
})
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -56,7 +56,7 @@ pub(super) fn is_empty_or_null_string(expr: &Expr) -> bool {
Expr::StringLiteral(ast::ExprStringLiteral { value, .. }) => value.is_empty(),
Expr::NoneLiteral(_) => true,
Expr::FString(ast::ExprFString { value, .. }) => {
value.parts().all(|f_string_part| match f_string_part {
value.iter().all(|f_string_part| match f_string_part {
ast::FStringPart::Literal(literal) => literal.is_empty(),
ast::FStringPart::FString(f_string) => f_string
.elements
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -70,7 +70,7 @@ pub(crate) fn assert_on_string_literal(checker: &mut Checker, test: &Expr) {
));
}
Expr::FString(ast::ExprFString { value, .. }) => {
let kind = if value.parts().all(|f_string_part| match f_string_part {
let kind = if value.iter().all(|f_string_part| match f_string_part {
ast::FStringPart::Literal(literal) => literal.is_empty(),
ast::FStringPart::FString(f_string) => {
f_string.elements.iter().all(|element| match element {
Expand All @@ -82,7 +82,7 @@ pub(crate) fn assert_on_string_literal(checker: &mut Checker, test: &Expr) {
}
}) {
Kind::Empty
} else if value.parts().any(|f_string_part| match f_string_part {
} else if value.iter().any(|f_string_part| match f_string_part {
ast::FStringPart::Literal(literal) => !literal.is_empty(),
ast::FStringPart::FString(f_string) => {
f_string.elements.iter().any(|element| match element {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -90,7 +90,7 @@ pub(crate) fn raise_vanilla_args(checker: &mut Checker, expr: &Expr) {
fn contains_message(expr: &Expr) -> bool {
match expr {
Expr::FString(ast::ExprFString { value, .. }) => {
for f_string_part in value.parts() {
for f_string_part in value {
match f_string_part {
ast::FStringPart::Literal(literal) => {
if literal.chars().any(char::is_whitespace) {
Expand Down
10 changes: 5 additions & 5 deletions crates/ruff_python_ast/src/comparable.rs
Original file line number Diff line number Diff line change
Expand Up @@ -583,10 +583,10 @@ impl<'a> From<ast::LiteralExpressionRef<'a>> for ComparableLiteral<'a> {
value, ..
}) => Self::Bool(value),
ast::LiteralExpressionRef::StringLiteral(ast::ExprStringLiteral { value, .. }) => {
Self::Str(value.parts().map(Into::into).collect())
Self::Str(value.iter().map(Into::into).collect())
}
ast::LiteralExpressionRef::BytesLiteral(ast::ExprBytesLiteral { value, .. }) => {
Self::Bytes(value.parts().map(Into::into).collect())
Self::Bytes(value.iter().map(Into::into).collect())
}
ast::LiteralExpressionRef::NumberLiteral(ast::ExprNumberLiteral { value, .. }) => {
Self::Number(value.into())
Expand Down Expand Up @@ -1012,17 +1012,17 @@ impl<'a> From<&'a ast::Expr> for ComparableExpr<'a> {
}),
ast::Expr::FString(ast::ExprFString { value, range: _ }) => {
Self::FString(ExprFString {
parts: value.parts().map(Into::into).collect(),
parts: value.iter().map(Into::into).collect(),
})
}
ast::Expr::StringLiteral(ast::ExprStringLiteral { value, range: _ }) => {
Self::StringLiteral(ExprStringLiteral {
parts: value.parts().map(Into::into).collect(),
parts: value.iter().map(Into::into).collect(),
})
}
ast::Expr::BytesLiteral(ast::ExprBytesLiteral { value, range: _ }) => {
Self::BytesLiteral(ExprBytesLiteral {
parts: value.parts().map(Into::into).collect(),
parts: value.iter().map(Into::into).collect(),
})
}
ast::Expr::NumberLiteral(ast::ExprNumberLiteral { value, range: _ }) => {
Expand Down
2 changes: 1 addition & 1 deletion crates/ruff_python_ast/src/helpers.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1326,7 +1326,7 @@ impl Truthiness {
Expr::NoneLiteral(_) => Self::Falsey,
Expr::EllipsisLiteral(_) => Self::Truthy,
Expr::FString(ast::ExprFString { value, .. }) => {
if value.parts().all(|part| match part {
if value.iter().all(|part| match part {
ast::FStringPart::Literal(string_literal) => string_literal.is_empty(),
ast::FStringPart::FString(f_string) => f_string.elements.is_empty(),
}) {
Expand Down
6 changes: 3 additions & 3 deletions crates/ruff_python_ast/src/node.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2742,7 +2742,7 @@ impl AstNode for ast::ExprFString {
{
let ast::ExprFString { value, range: _ } = self;

for f_string_part in value.parts() {
for f_string_part in value {
match f_string_part {
ast::FStringPart::Literal(string_literal) => {
visitor.visit_string_literal(string_literal);
Expand Down Expand Up @@ -2788,7 +2788,7 @@ impl AstNode for ast::ExprStringLiteral {
{
let ast::ExprStringLiteral { value, range: _ } = self;

for string_literal in value.parts() {
for string_literal in value {
visitor.visit_string_literal(string_literal);
}
}
Expand Down Expand Up @@ -2827,7 +2827,7 @@ impl AstNode for ast::ExprBytesLiteral {
{
let ast::ExprBytesLiteral { value, range: _ } = self;

for bytes_literal in value.parts() {
for bytes_literal in value {
visitor.visit_bytes_literal(bytes_literal);
}
}
Expand Down
130 changes: 93 additions & 37 deletions crates/ruff_python_ast/src/nodes.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,8 +4,8 @@ use std::cell::OnceCell;
use std::fmt;
use std::fmt::Debug;
use std::ops::Deref;
use std::slice::{Iter, IterMut};

use itertools::Either::{Left, Right};
use itertools::Itertools;

use ruff_text_size::{Ranged, TextRange, TextSize};
Expand Down Expand Up @@ -1051,23 +1051,33 @@ impl FStringValue {
matches!(self.inner, FStringValueInner::Concatenated(_))
}

/// Returns an iterator over all the [`FStringPart`]s contained in this value.
pub fn parts(&self) -> impl Iterator<Item = &FStringPart> {
/// Returns a slice of all the [`FStringPart`]s contained in this value.
pub fn as_slice(&self) -> &[FStringPart] {
match &self.inner {
FStringValueInner::Single(part) => Left(std::iter::once(part)),
FStringValueInner::Concatenated(parts) => Right(parts.iter()),
FStringValueInner::Single(part) => std::slice::from_ref(part),
FStringValueInner::Concatenated(parts) => parts,
}
}

/// Returns an iterator over all the [`FStringPart`]s contained in this value
/// that allows modification.
pub(crate) fn parts_mut(&mut self) -> impl Iterator<Item = &mut FStringPart> {
/// Returns a mutable slice of all the [`FStringPart`]s contained in this value.
fn as_mut_slice(&mut self) -> &mut [FStringPart] {
match &mut self.inner {
FStringValueInner::Single(part) => Left(std::iter::once(part)),
FStringValueInner::Concatenated(parts) => Right(parts.iter_mut()),
FStringValueInner::Single(part) => std::slice::from_mut(part),
FStringValueInner::Concatenated(parts) => parts,
}
}

/// Returns an iterator over all the [`FStringPart`]s contained in this value.
pub fn iter(&self) -> Iter<FStringPart> {
self.as_slice().iter()
}

/// Returns an iterator over all the [`FStringPart`]s contained in this value
/// that allows modification.
pub(crate) fn iter_mut(&mut self) -> IterMut<FStringPart> {
self.as_mut_slice().iter_mut()
}

/// Returns an iterator over the [`StringLiteral`] parts contained in this value.
///
/// Note that this doesn't nest into the f-string parts. For example,
Expand All @@ -1078,7 +1088,7 @@ impl FStringValue {
///
/// Here, the string literal parts returned would be `"foo"` and `"baz"`.
pub fn literals(&self) -> impl Iterator<Item = &StringLiteral> {
self.parts().filter_map(|part| part.as_literal())
self.iter().filter_map(|part| part.as_literal())
}

/// Returns an iterator over the [`FString`] parts contained in this value.
Expand All @@ -1091,7 +1101,7 @@ impl FStringValue {
///
/// Here, the f-string parts returned would be `f"bar {x}"` and `f"qux"`.
pub fn f_strings(&self) -> impl Iterator<Item = &FString> {
self.parts().filter_map(|part| part.as_f_string())
self.iter().filter_map(|part| part.as_f_string())
}

/// Returns an iterator over all the [`FStringElement`] contained in this value.
Expand All @@ -1110,6 +1120,15 @@ impl FStringValue {
}
}

impl<'a> IntoIterator for &'a FStringValue {
type Item = &'a FStringPart;
type IntoIter = Iter<'a, FStringPart>;

fn into_iter(self) -> Self::IntoIter {
self.iter()
}
}

/// An internal representation of [`FStringValue`].
#[derive(Clone, Debug, PartialEq)]
enum FStringValueInner {
Expand Down Expand Up @@ -1238,26 +1257,36 @@ impl StringLiteralValue {
/// For an implicitly concatenated string, it returns `true` only if the first
/// string literal is a unicode string.
pub fn is_unicode(&self) -> bool {
self.parts().next().map_or(false, |part| part.unicode)
self.iter().next().map_or(false, |part| part.unicode)
}

/// Returns an iterator over all the [`StringLiteral`] parts contained in this value.
pub fn parts(&self) -> impl Iterator<Item = &StringLiteral> {
/// Returns a slice of all the [`StringLiteral`] parts contained in this value.
pub fn as_slice(&self) -> &[StringLiteral] {
match &self.inner {
StringLiteralValueInner::Single(value) => Left(std::iter::once(value)),
StringLiteralValueInner::Concatenated(value) => Right(value.strings.iter()),
StringLiteralValueInner::Single(value) => std::slice::from_ref(value),
StringLiteralValueInner::Concatenated(value) => value.strings.as_slice(),
}
}

/// Returns an iterator over all the [`StringLiteral`] parts contained in this value
/// that allows modification.
pub(crate) fn parts_mut(&mut self) -> impl Iterator<Item = &mut StringLiteral> {
/// Returns a mutable slice of all the [`StringLiteral`] parts contained in this value.
fn as_mut_slice(&mut self) -> &mut [StringLiteral] {
match &mut self.inner {
StringLiteralValueInner::Single(value) => Left(std::iter::once(value)),
StringLiteralValueInner::Concatenated(value) => Right(value.strings.iter_mut()),
StringLiteralValueInner::Single(value) => std::slice::from_mut(value),
StringLiteralValueInner::Concatenated(value) => value.strings.as_mut_slice(),
}
}

/// Returns an iterator over all the [`StringLiteral`] parts contained in this value.
pub fn iter(&self) -> Iter<StringLiteral> {
self.as_slice().iter()
}

/// Returns an iterator over all the [`StringLiteral`] parts contained in this value
/// that allows modification.
pub(crate) fn iter_mut(&mut self) -> IterMut<StringLiteral> {
self.as_mut_slice().iter_mut()
}

/// Returns `true` if the string literal value is empty.
pub fn is_empty(&self) -> bool {
self.len() == 0
Expand All @@ -1266,12 +1295,12 @@ impl StringLiteralValue {
/// Returns the total length of the string literal value, in bytes, not
/// [`char`]s or graphemes.
pub fn len(&self) -> usize {
self.parts().fold(0, |acc, part| acc + part.value.len())
self.iter().fold(0, |acc, part| acc + part.value.len())
}

/// Returns an iterator over the [`char`]s of each string literal part.
pub fn chars(&self) -> impl Iterator<Item = char> + '_ {
self.parts().flat_map(|part| part.value.chars())
self.iter().flat_map(|part| part.value.chars())
}

/// Returns the concatenated string value as a [`str`].
Expand All @@ -1286,6 +1315,15 @@ impl StringLiteralValue {
}
}

impl<'a> IntoIterator for &'a StringLiteralValue {
type Item = &'a StringLiteral;
type IntoIter = Iter<'a, StringLiteral>;

fn into_iter(self) -> Self::IntoIter {
self.iter()
}
}

impl PartialEq<str> for StringLiteralValue {
fn eq(&self, other: &str) -> bool {
if self.len() != other.len() {
Expand Down Expand Up @@ -1457,37 +1495,55 @@ impl BytesLiteralValue {
matches!(self.inner, BytesLiteralValueInner::Concatenated(_))
}

/// Returns an iterator over all the [`BytesLiteral`] parts contained in this value.
pub fn parts(&self) -> impl Iterator<Item = &BytesLiteral> {
/// Returns a slice of all the [`BytesLiteral`] parts contained in this value.
pub fn as_slice(&self) -> &[BytesLiteral] {
match &self.inner {
BytesLiteralValueInner::Single(value) => Left(std::iter::once(value)),
BytesLiteralValueInner::Concatenated(values) => Right(values.iter()),
BytesLiteralValueInner::Single(value) => std::slice::from_ref(value),
BytesLiteralValueInner::Concatenated(value) => value.as_slice(),
}
}

/// Returns an iterator over all the [`BytesLiteral`] parts contained in this value
/// that allows modification.
pub(crate) fn parts_mut(&mut self) -> impl Iterator<Item = &mut BytesLiteral> {
/// Returns a mutable slice of all the [`BytesLiteral`] parts contained in this value.
fn as_mut_slice(&mut self) -> &mut [BytesLiteral] {
match &mut self.inner {
BytesLiteralValueInner::Single(value) => Left(std::iter::once(value)),
BytesLiteralValueInner::Concatenated(values) => Right(values.iter_mut()),
BytesLiteralValueInner::Single(value) => std::slice::from_mut(value),
BytesLiteralValueInner::Concatenated(value) => value.as_mut_slice(),
}
}

/// Returns an iterator over all the [`BytesLiteral`] parts contained in this value.
pub fn iter(&self) -> Iter<BytesLiteral> {
self.as_slice().iter()
}

/// Returns an iterator over all the [`BytesLiteral`] parts contained in this value
/// that allows modification.
pub(crate) fn iter_mut(&mut self) -> IterMut<BytesLiteral> {
self.as_mut_slice().iter_mut()
}

/// Returns `true` if the concatenated bytes has a length of zero.
pub fn is_empty(&self) -> bool {
self.parts().all(|part| part.is_empty())
self.iter().all(|part| part.is_empty())
}

/// Returns the length of the concatenated bytes.
pub fn len(&self) -> usize {
self.parts().map(|part| part.len()).sum()
self.iter().map(|part| part.len()).sum()
}

/// Returns an iterator over the bytes of the concatenated bytes.
fn bytes(&self) -> impl Iterator<Item = u8> + '_ {
self.parts()
.flat_map(|part| part.as_slice().iter().copied())
self.iter().flat_map(|part| part.as_slice().iter().copied())
}
}

impl<'a> IntoIterator for &'a BytesLiteralValue {
type Item = &'a BytesLiteral;
type IntoIter = Iter<'a, BytesLiteral>;

fn into_iter(self) -> Self::IntoIter {
self.iter()
}
}

Expand Down
6 changes: 3 additions & 3 deletions crates/ruff_python_ast/src/visitor.rs
Original file line number Diff line number Diff line change
Expand Up @@ -477,7 +477,7 @@ pub fn walk_expr<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, expr: &'a Expr) {
visitor.visit_arguments(arguments);
}
Expr::FString(ast::ExprFString { value, .. }) => {
for part in value.parts() {
for part in value {
match part {
FStringPart::Literal(string_literal) => {
visitor.visit_string_literal(string_literal);
Expand All @@ -487,12 +487,12 @@ pub fn walk_expr<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, expr: &'a Expr) {
}
}
Expr::StringLiteral(ast::ExprStringLiteral { value, .. }) => {
for string_literal in value.parts() {
for string_literal in value {
visitor.visit_string_literal(string_literal);
}
}
Expr::BytesLiteral(ast::ExprBytesLiteral { value, .. }) => {
for bytes_literal in value.parts() {
for bytes_literal in value {
visitor.visit_bytes_literal(bytes_literal);
}
}
Expand Down
Loading

0 comments on commit 18452cf

Please sign in to comment.