Skip to main content

openjd_expr/functions/
string.rs

1// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2// Copyright by contributors to this project.
3// SPDX-License-Identifier: (Apache-2.0 OR MIT)
4
5//! String function implementations.
6
7use crate::error::ExpressionError;
8use crate::function_library::EvalContext;
9use crate::types::ExprType;
10use crate::value::ExprValue;
11
12type R = Result<ExprValue, ExpressionError>;
13type Ctx<'a> = &'a mut dyn EvalContext;
14
15fn get_str(a: &ExprValue) -> Result<&str, ExpressionError> {
16    match a {
17        ExprValue::String(s) => Ok(s),
18        ExprValue::Path { value, .. } => Ok(value),
19        _ => Err(ExpressionError::new(format!(
20            "String method not supported on {}",
21            a.expr_type()
22        ))),
23    }
24}
25
26/// Convert a byte offset (from str::find) to a codepoint offset (matching Python).
27fn byte_to_char_offset(s: &str, byte_offset: usize) -> usize {
28    s[..byte_offset].chars().count()
29}
30
31pub fn upper_fn(ctx: Ctx, a: &[ExprValue]) -> R {
32    let s = get_str(&a[0])?;
33    ctx.count_string_ops(s.len())?;
34    Ok(ExprValue::String(s.to_uppercase()))
35}
36pub fn lower_fn(ctx: Ctx, a: &[ExprValue]) -> R {
37    let s = get_str(&a[0])?;
38    ctx.count_string_ops(s.len())?;
39    Ok(ExprValue::String(s.to_lowercase()))
40}
41
42pub fn strip_fn(ctx: Ctx, a: &[ExprValue]) -> R {
43    let s = get_str(&a[0])?;
44    ctx.count_string_ops(s.len())?;
45    if a.len() > 1 {
46        let chars: Vec<char> = get_str(&a[1])?.chars().collect();
47        Ok(ExprValue::String(
48            s.trim_matches(|c| chars.contains(&c)).to_string(),
49        ))
50    } else {
51        Ok(ExprValue::String(s.trim().to_string()))
52    }
53}
54
55pub fn lstrip_fn(ctx: Ctx, a: &[ExprValue]) -> R {
56    let s = get_str(&a[0])?;
57    ctx.count_string_ops(s.len())?;
58    if a.len() > 1 {
59        let chars: Vec<char> = get_str(&a[1])?.chars().collect();
60        Ok(ExprValue::String(
61            s.trim_start_matches(|c| chars.contains(&c)).to_string(),
62        ))
63    } else {
64        Ok(ExprValue::String(s.trim_start().to_string()))
65    }
66}
67
68pub fn rstrip_fn(ctx: Ctx, a: &[ExprValue]) -> R {
69    let s = get_str(&a[0])?;
70    ctx.count_string_ops(s.len())?;
71    if a.len() > 1 {
72        let chars: Vec<char> = get_str(&a[1])?.chars().collect();
73        Ok(ExprValue::String(
74            s.trim_end_matches(|c| chars.contains(&c)).to_string(),
75        ))
76    } else {
77        Ok(ExprValue::String(s.trim_end().to_string()))
78    }
79}
80
81pub fn removeprefix_fn(ctx: Ctx, a: &[ExprValue]) -> R {
82    let s = get_str(&a[0])?;
83    let prefix = get_str(&a[1])?;
84    ctx.count_string_ops(s.len())?;
85    Ok(ExprValue::String(
86        s.strip_prefix(prefix).unwrap_or(s).to_string(),
87    ))
88}
89
90pub fn removesuffix_fn(ctx: Ctx, a: &[ExprValue]) -> R {
91    let s = get_str(&a[0])?;
92    let suffix = get_str(&a[1])?;
93    ctx.count_string_ops(s.len())?;
94    Ok(ExprValue::String(
95        s.strip_suffix(suffix).unwrap_or(s).to_string(),
96    ))
97}
98
99pub fn replace_fn(ctx: Ctx, a: &[ExprValue]) -> R {
100    let s = get_str(&a[0])?;
101    let old = get_str(&a[1])?;
102    let new = get_str(&a[2])?;
103    ctx.count_string_ops(s.len())?;
104    if old.is_empty() {
105        return Err(ExpressionError::new("replace failed: empty old string"));
106    }
107    Ok(ExprValue::String(s.replace(old, new)))
108}
109
110pub fn startswith_fn(ctx: Ctx, a: &[ExprValue]) -> R {
111    let s = get_str(&a[0])?;
112    ctx.count_string_ops(s.len())?;
113    Ok(ExprValue::Bool(s.starts_with(get_str(&a[1])?)))
114}
115
116pub fn endswith_fn(ctx: Ctx, a: &[ExprValue]) -> R {
117    let s = get_str(&a[0])?;
118    ctx.count_string_ops(s.len())?;
119    Ok(ExprValue::Bool(s.ends_with(get_str(&a[1])?)))
120}
121
122pub fn find_fn(ctx: Ctx, a: &[ExprValue]) -> R {
123    let s = get_str(&a[0])?;
124    let sub = get_str(&a[1])?;
125    if sub.is_empty() {
126        return Err(ExpressionError::new("find failed: empty substring"));
127    }
128    ctx.count_string_ops(s.len())?;
129    Ok(ExprValue::Int(
130        s.find(sub)
131            .map(|p| byte_to_char_offset(s, p) as i64)
132            .unwrap_or(-1),
133    ))
134}
135
136pub fn rfind_fn(ctx: Ctx, a: &[ExprValue]) -> R {
137    let s = get_str(&a[0])?;
138    let sub = get_str(&a[1])?;
139    if sub.is_empty() {
140        return Err(ExpressionError::new("rfind failed: empty substring"));
141    }
142    ctx.count_string_ops(s.len())?;
143    Ok(ExprValue::Int(
144        s.rfind(sub)
145            .map(|p| byte_to_char_offset(s, p) as i64)
146            .unwrap_or(-1),
147    ))
148}
149
150pub fn index_fn(ctx: Ctx, a: &[ExprValue]) -> R {
151    let s = get_str(&a[0])?;
152    let sub = get_str(&a[1])?;
153    if sub.is_empty() {
154        return Err(ExpressionError::new("index failed: empty substring"));
155    }
156    ctx.count_string_ops(s.len())?;
157    match s.find(sub) {
158        Some(p) => Ok(ExprValue::Int(byte_to_char_offset(s, p) as i64)),
159        None => Err(ExpressionError::new(format!(
160            "index failed: substring '{sub}' not found"
161        ))),
162    }
163}
164
165pub fn rindex_fn(ctx: Ctx, a: &[ExprValue]) -> R {
166    let s = get_str(&a[0])?;
167    let sub = get_str(&a[1])?;
168    if sub.is_empty() {
169        return Err(ExpressionError::new("rindex failed: empty substring"));
170    }
171    ctx.count_string_ops(s.len())?;
172    match s.rfind(sub) {
173        Some(p) => Ok(ExprValue::Int(byte_to_char_offset(s, p) as i64)),
174        None => Err(ExpressionError::new(format!(
175            "rindex failed: substring '{sub}' not found"
176        ))),
177    }
178}
179
180pub fn count_fn(ctx: Ctx, a: &[ExprValue]) -> R {
181    let s = get_str(&a[0])?;
182    let sub = get_str(&a[1])?;
183    if sub.is_empty() {
184        return Err(ExpressionError::new("count failed: empty substring"));
185    }
186    ctx.count_string_ops(s.len())?;
187    Ok(ExprValue::Int(s.matches(sub).count() as i64))
188}
189
190pub fn split_fn(ctx: Ctx, a: &[ExprValue]) -> R {
191    let s = get_str(&a[0])?;
192    if a.len() == 1 {
193        // Whitespace split
194        ctx.count_string_ops(s.len())?;
195        let parts: Vec<ExprValue> = s
196            .split_whitespace()
197            .map(|p| ExprValue::String(p.to_string()))
198            .collect();
199        return ExprValue::make_list_checked(ctx, parts, ExprType::STRING);
200    }
201    let sep = get_str(&a[1])?;
202    if sep.is_empty() {
203        return Err(ExpressionError::new("split failed: empty separator"));
204    }
205    ctx.count_string_ops(s.len())?;
206    let maxsplit = a.get(2).and_then(|v| match v {
207        ExprValue::Int(n) => Some(*n as usize),
208        _ => None,
209    });
210    let parts: Vec<ExprValue> = match maxsplit {
211        Some(n) => s
212            .splitn(n + 1, sep)
213            .map(|p| ExprValue::String(p.to_string()))
214            .collect(),
215        None => s
216            .split(sep)
217            .map(|p| ExprValue::String(p.to_string()))
218            .collect(),
219    };
220    ExprValue::make_list_checked(ctx, parts, ExprType::STRING)
221}
222
223pub fn rsplit_fn(ctx: Ctx, a: &[ExprValue]) -> R {
224    let s = get_str(&a[0])?;
225    if a.len() == 1 {
226        ctx.count_string_ops(s.len())?;
227        let parts: Vec<ExprValue> = s
228            .split_whitespace()
229            .map(|p| ExprValue::String(p.to_string()))
230            .collect();
231        return ExprValue::make_list_checked(ctx, parts, ExprType::STRING);
232    }
233    let sep = get_str(&a[1])?;
234    if sep.is_empty() {
235        return Err(ExpressionError::new("split failed: empty separator"));
236    }
237    ctx.count_string_ops(s.len())?;
238    let maxsplit = a.get(2).and_then(|v| match v {
239        ExprValue::Int(n) => Some(*n as usize),
240        _ => None,
241    });
242    let parts: Vec<ExprValue> = match maxsplit {
243        Some(n) => {
244            let mut v: Vec<_> = s
245                .rsplitn(n + 1, sep)
246                .map(|p| ExprValue::String(p.to_string()))
247                .collect();
248            v.reverse();
249            v
250        }
251        None => s
252            .split(sep)
253            .map(|p| ExprValue::String(p.to_string()))
254            .collect(),
255    };
256    ExprValue::make_list_checked(ctx, parts, ExprType::STRING)
257}
258
259pub fn isdigit_fn(ctx: Ctx, a: &[ExprValue]) -> R {
260    let s = get_str(&a[0])?;
261    ctx.count_string_ops(s.len())?;
262    Ok(ExprValue::Bool(
263        !s.is_empty() && s.chars().all(|c| c.is_ascii_digit()),
264    ))
265}
266pub fn isalpha_fn(ctx: Ctx, a: &[ExprValue]) -> R {
267    let s = get_str(&a[0])?;
268    ctx.count_string_ops(s.len())?;
269    Ok(ExprValue::Bool(
270        !s.is_empty() && s.chars().all(|c| c.is_alphabetic()),
271    ))
272}
273pub fn isalnum_fn(ctx: Ctx, a: &[ExprValue]) -> R {
274    let s = get_str(&a[0])?;
275    ctx.count_string_ops(s.len())?;
276    Ok(ExprValue::Bool(
277        !s.is_empty() && s.chars().all(|c| c.is_alphanumeric()),
278    ))
279}
280pub fn isspace_fn(ctx: Ctx, a: &[ExprValue]) -> R {
281    let s = get_str(&a[0])?;
282    ctx.count_string_ops(s.len())?;
283    Ok(ExprValue::Bool(
284        !s.is_empty() && s.chars().all(|c| c.is_whitespace()),
285    ))
286}
287pub fn isupper_fn(ctx: Ctx, a: &[ExprValue]) -> R {
288    let s = get_str(&a[0])?;
289    ctx.count_string_ops(s.len())?;
290    Ok(ExprValue::Bool(
291        s.chars().any(|c| c.is_alphabetic())
292            && s.chars()
293                .filter(|c| c.is_alphabetic())
294                .all(|c| c.is_uppercase()),
295    ))
296}
297pub fn islower_fn(ctx: Ctx, a: &[ExprValue]) -> R {
298    let s = get_str(&a[0])?;
299    ctx.count_string_ops(s.len())?;
300    Ok(ExprValue::Bool(
301        s.chars().any(|c| c.is_alphabetic())
302            && s.chars()
303                .filter(|c| c.is_alphabetic())
304                .all(|c| c.is_lowercase()),
305    ))
306}
307pub fn isascii_fn(ctx: Ctx, a: &[ExprValue]) -> R {
308    let s = get_str(&a[0])?;
309    ctx.count_string_ops(s.len())?;
310    Ok(ExprValue::Bool(s.is_ascii()))
311}
312
313pub fn title_fn(ctx: Ctx, a: &[ExprValue]) -> R {
314    let s = get_str(&a[0])?;
315    ctx.count_string_ops(s.len())?;
316    let mut result = String::with_capacity(s.len());
317    let mut capitalize_next = true;
318    for c in s.chars() {
319        if c.is_alphanumeric() {
320            if capitalize_next {
321                result.extend(c.to_uppercase());
322                capitalize_next = false;
323            } else {
324                result.extend(c.to_lowercase());
325            }
326        } else {
327            result.push(c);
328            capitalize_next = true;
329        }
330    }
331    Ok(ExprValue::String(result))
332}
333
334pub fn capitalize_fn(ctx: Ctx, a: &[ExprValue]) -> R {
335    let s = get_str(&a[0])?;
336    ctx.count_string_ops(s.len())?;
337    let mut chars = s.chars();
338    let result = match chars.next() {
339        None => String::new(),
340        Some(c) => c.to_uppercase().to_string() + &chars.as_str().to_lowercase(),
341    };
342    Ok(ExprValue::String(result))
343}
344
345pub fn center_fn(ctx: Ctx, a: &[ExprValue]) -> R {
346    let s = get_str(&a[0])?;
347    let width = match &a[1] {
348        ExprValue::Int(w) => *w as usize,
349        _ => return Err(ExpressionError::new("center() width must be int")),
350    };
351    ctx.count_string_ops(width.max(s.len()))?;
352    let clen = s.chars().count();
353    if clen >= width {
354        return Ok(ExprValue::String(s.to_string()));
355    }
356    let pad = width - clen;
357    let left = pad / 2;
358    let right = pad - left;
359    Ok(ExprValue::String(format!(
360        "{}{}{}",
361        " ".repeat(left),
362        s,
363        " ".repeat(right)
364    )))
365}
366
367pub fn ljust_fn(ctx: Ctx, a: &[ExprValue]) -> R {
368    let s = get_str(&a[0])?;
369    let width = match &a[1] {
370        ExprValue::Int(w) => *w as usize,
371        _ => return Err(ExpressionError::new("ljust() width must be int")),
372    };
373    ctx.count_string_ops(width.max(s.len()))?;
374    Ok(ExprValue::String(format!("{:<width$}", s, width = width)))
375}
376
377pub fn rjust_fn(ctx: Ctx, a: &[ExprValue]) -> R {
378    let s = get_str(&a[0])?;
379    let width = match &a[1] {
380        ExprValue::Int(w) => *w as usize,
381        _ => return Err(ExpressionError::new("rjust() width must be int")),
382    };
383    ctx.count_string_ops(width.max(s.len()))?;
384    Ok(ExprValue::String(format!("{:>width$}", s, width = width)))
385}