use indexmap::IndexMap;
use regex::Regex;
use crate::{
error::{EvalError, EvalResult, InterpreterError},
eval::modules::arg_str,
value::{ExceptionValue, MatchGroup, MatchValue, Value, ValueKey, shared_list},
};
pub fn has_function(name: &str) -> bool {
matches!(name, "findall" | "sub" | "split" | "match" | "search" | "fullmatch")
}
pub fn call(func: &str, args: &[Value], kwargs: &IndexMap<String, Value>) -> EvalResult {
match func {
"findall" => findall(args),
"sub" => sub(func, args, kwargs),
"split" => split(func, args, kwargs),
"match" => {
anchored_search(func, args, true, false)
}
"fullmatch" => anchored_search(func, args, true, true),
"search" => anchored_search(func, args, false, false),
_ => {
Err(InterpreterError::AttributeError(format!("module 're' has no attribute '{func}'"))
.into())
}
}
}
fn count_arg(args: &[Value], pos: usize, kwargs: &IndexMap<String, Value>, key: &str) -> usize {
let value = args.get(pos).or_else(|| kwargs.get(key));
match value {
Some(Value::Int(n)) if *n > 0 => usize::try_from(*n).unwrap_or(0),
_ => 0,
}
}
fn compile(pattern: &str) -> Result<Regex, EvalError> {
Regex::new(pattern)
.map_err(|e| EvalError::Exception(ExceptionValue::new("error", format!("{e}"))))
}
fn findall(args: &[Value]) -> EvalResult {
let pattern = arg_str("findall", args, 0)?;
let text = arg_str("findall", args, 1)?;
let re = compile(pattern)?;
let group_count = re.captures_len().saturating_sub(1);
let mut result = Vec::new();
if group_count == 0 {
for m in re.find_iter(text) {
result.push(Value::String(m.as_str().into()));
}
} else {
for caps in re.captures_iter(text) {
if group_count == 1 {
result.push(Value::String(group_text(&caps, 1)));
} else {
let groups =
(1..=group_count).map(|i| Value::String(group_text(&caps, i))).collect();
result.push(Value::Tuple(groups));
}
}
}
Ok(Value::List(shared_list(result)))
}
fn sub(func: &str, args: &[Value], kwargs: &IndexMap<String, Value>) -> EvalResult {
let pattern = arg_str(func, args, 0)?;
let repl = arg_str(func, args, 1)?;
let text = arg_str(func, args, 2)?;
let count = count_arg(args, 3, kwargs, "count");
let re = compile(pattern)?;
let translated = translate_python_repl(repl);
let replaced = if count == 0 {
re.replace_all(text, translated.as_str())
} else {
re.replacen(text, count, translated.as_str())
};
Ok(Value::String(replaced.into_owned().into()))
}
fn translate_python_repl(repl: &str) -> String {
use std::fmt::Write as _;
let mut out = String::with_capacity(repl.len());
let mut chars = repl.chars().peekable();
while let Some(c) = chars.next() {
match c {
'$' => out.push_str("$$"),
'\\' => match chars.peek() {
Some(&n) if n.is_ascii_digit() => {
chars.next();
let _ = write!(out, "${{{n}}}");
}
Some(&'g') => {
chars.next();
if matches!(chars.peek(), Some(&'<')) {
chars.next();
let mut name = String::new();
while let Some(&ch) = chars.peek() {
if ch == '>' {
chars.next();
break;
}
name.push(ch);
chars.next();
}
let _ = write!(out, "${{{name}}}");
} else {
out.push_str("\\g");
}
}
Some(&'\\') => {
chars.next();
out.push('\\');
}
Some(&'n') => {
chars.next();
out.push('\n');
}
Some(&'t') => {
chars.next();
out.push('\t');
}
_ => out.push('\\'),
},
other => out.push(other),
}
}
out
}
fn split(func: &str, args: &[Value], kwargs: &IndexMap<String, Value>) -> EvalResult {
let pattern = arg_str(func, args, 0)?;
let text = arg_str(func, args, 1)?;
let maxsplit = count_arg(args, 2, kwargs, "maxsplit");
let re = compile(pattern)?;
let parts: Vec<Value> = if maxsplit == 0 {
re.split(text).map(|p| Value::String(p.into())).collect()
} else {
re.splitn(text, maxsplit + 1).map(|p| Value::String(p.into())).collect()
};
Ok(Value::List(shared_list(parts)))
}
fn anchored_search(func: &str, args: &[Value], anchor_start: bool, anchor_end: bool) -> EvalResult {
let pattern = arg_str(func, args, 0)?;
let text = arg_str(func, args, 1)?;
let re = compile(pattern)?;
let Some(caps) = re.captures(text) else {
return Ok(Value::None);
};
let Some(whole) = caps.get(0) else {
return Ok(Value::None);
};
if (anchor_start && whole.start() != 0) || (anchor_end && whole.end() != text.len()) {
return Ok(Value::None);
}
Ok(Value::ReMatch(Box::new(build_match(&caps, &re, text))))
}
fn build_match(caps: ®ex::Captures<'_>, re: &Regex, text: &str) -> MatchValue {
let groups = (0..caps.len())
.map(|i| {
caps.get(i).map(|m| MatchGroup {
text: m.as_str().to_string(),
start: char_offset(text, m.start()),
end: char_offset(text, m.end()),
})
})
.collect();
let mut named = indexmap::IndexMap::new();
for (index, name) in re.capture_names().enumerate() {
if let Some(name) = name {
named.insert(name.to_string(), index);
}
}
MatchValue { groups, named }
}
fn char_offset(text: &str, byte: usize) -> usize {
text.get(..byte).map_or(byte, |prefix| prefix.chars().count())
}
fn group_text(caps: ®ex::Captures<'_>, index: usize) -> compact_str::CompactString {
caps.get(index).map_or_else(compact_str::CompactString::default, |m| m.as_str().into())
}
pub fn dispatch_match_method(
m: &MatchValue,
method: &str,
args: &[Value],
kwargs: &indexmap::IndexMap<String, Value>,
) -> EvalResult {
crate::eval::functions::reject_kwargs(method, kwargs)?;
match method {
"group" => match args.len() {
0 => group_value(m, 0),
1 => group_by_arg(m, &args[0]),
_ => {
let mut out = Vec::with_capacity(args.len());
for arg in args {
out.push(group_by_arg(m, arg)?);
}
Ok(Value::Tuple(out))
}
},
"groups" => {
let default = args.first().cloned().unwrap_or(Value::None);
let out =
m.groups.iter().skip(1).map(|g| group_or_default(g.as_ref(), &default)).collect();
Ok(Value::Tuple(out))
}
"groupdict" => {
let mut map = IndexMap::new();
for (name, &index) in &m.named {
let value = m
.groups
.get(index)
.and_then(Option::as_ref)
.map_or(Value::None, |g| Value::String(g.text.as_str().into()));
map.insert(ValueKey::String(name.as_str().into()), value);
}
Ok(Value::Dict(map))
}
"start" => Ok(Value::Int(group_span(m, args)?.0)),
"end" => Ok(Value::Int(group_span(m, args)?.1)),
"span" => {
let (start, end) = group_span(m, args)?;
Ok(Value::Tuple(vec![Value::Int(start), Value::Int(end)]))
}
_ => Err(InterpreterError::AttributeError(format!(
"'re.Match' object has no attribute '{method}'"
))
.into()),
}
}
fn group_by_arg(m: &MatchValue, arg: &Value) -> EvalResult {
match arg {
Value::Int(i) => group_value(m, group_index(*i)?),
Value::String(name) => {
let index = *m.named.get(name.as_str()).ok_or_else(|| no_such_group(name.as_str()))?;
group_value(m, index)
}
other => Err(InterpreterError::TypeError(format!(
"group indices must be integers or strings, not '{}'",
other.type_name()
))
.into()),
}
}
fn group_value(m: &MatchValue, index: usize) -> EvalResult {
let group = m.groups.get(index).ok_or_else(|| no_such_group(&index.to_string()))?;
Ok(group_or_default(group.as_ref(), &Value::None))
}
fn group_span(m: &MatchValue, args: &[Value]) -> Result<(i64, i64), EvalError> {
let index = match args.first() {
None => 0,
Some(Value::Int(i)) => group_index(*i)?,
Some(Value::String(name)) => {
*m.named.get(name.as_str()).ok_or_else(|| no_such_group(name.as_str()))?
}
Some(other) => {
return Err(InterpreterError::TypeError(format!(
"group indices must be integers or strings, not '{}'",
other.type_name()
))
.into());
}
};
let group = m.groups.get(index).ok_or_else(|| no_such_group(&index.to_string()))?;
Ok(group.as_ref().map_or((-1, -1), |g| {
(i64::try_from(g.start).unwrap_or(-1), i64::try_from(g.end).unwrap_or(-1))
}))
}
fn group_or_default(group: Option<&MatchGroup>, default: &Value) -> Value {
group.map_or_else(|| default.clone(), |g| Value::String(g.text.as_str().into()))
}
fn group_index(i: i64) -> Result<usize, EvalError> {
usize::try_from(i).map_err(|_| no_such_group(&i.to_string()))
}
fn no_such_group(name: &str) -> EvalError {
EvalError::Exception(ExceptionValue::new("IndexError", format!("no such group: {name}")))
}
pub struct ReModule;
#[async_trait::async_trait]
impl crate::eval::modules::Module for ReModule {
fn name(&self) -> &'static str {
"re"
}
fn has_function(&self, name: &str) -> bool {
has_function(name)
}
async fn call(
&self,
_state: &mut crate::state::InterpreterState,
func: &str,
args: &[Value],
kwargs: &IndexMap<String, Value>,
_tools: &crate::tools::Tools,
) -> EvalResult {
call(func, args, kwargs)
}
}