use anyhow::anyhow;
use crate::{
self as starlark,
environment::MethodsBuilder,
stdlib::util::{convert_index, convert_indices},
values::{
list::{List, ListRef},
none::{NoneOr, NoneType},
Value, ValueError,
},
};
#[starlark_module]
pub(crate) fn list_methods(builder: &mut MethodsBuilder) {
fn append(this: Value, ref el: Value) -> anyhow::Result<NoneType> {
let this = List::from_value_mut(this)?;
this.push(el, heap);
Ok(NoneType)
}
fn clear(this: Value) -> anyhow::Result<NoneType> {
let this = List::from_value_mut(this)?;
this.clear();
Ok(NoneType)
}
fn extend(this: Value, ref other: Value) -> anyhow::Result<NoneType> {
let res = List::from_value_mut(this)?;
if this.ptr_eq(other) {
res.double(heap);
} else {
other.with_iterator(heap, |it| {
res.extend(it, heap);
})?;
}
Ok(NoneType)
}
#[starlark(speculative_exec_safe)]
fn index(
this: &ListRef,
ref needle: Value,
ref start @ NoneOr::None: NoneOr<i32>,
ref end @ NoneOr::None: NoneOr<i32>,
) -> anyhow::Result<i32> {
let (start, end) = convert_indices(this.len() as i32, start, end);
if let Some(haystack) = this.get(start..end) {
for (i, x) in haystack.iter().enumerate() {
if x.equals(needle)? {
return Ok((i + start) as i32);
}
}
}
Err(anyhow!("Element '{}' not found in '{}'", needle, this))
}
fn insert(this: Value, ref index: i32, ref el: Value) -> anyhow::Result<NoneType> {
let this = List::from_value_mut(this)?;
let index = convert_index(this.len() as i32, index);
this.insert(index, el, heap);
Ok(NoneType)
}
fn pop<'v>(this: Value, ref index: Option<Value>) -> anyhow::Result<Value<'v>> {
let index = match index {
Some(index) => Some(index.to_int()?),
None => None,
};
let this = List::from_value_mut(this)?;
let index = index.unwrap_or_else(|| (this.len() as i32) - 1);
if index < 0 || index >= this.len() as i32 {
return Err(ValueError::IndexOutOfBound(index).into());
}
Ok(this.remove(index as usize))
}
fn remove(this: Value, ref needle: Value) -> anyhow::Result<NoneType> {
let position = {
let this = List::from_value(this).unwrap();
let position = this.iter().position(|v| v == needle);
match position {
Some(i) => i,
None => {
return Err(anyhow!("Element '{}' not found in list '{}'", needle, this));
}
}
};
{
let this = List::from_value_mut(this)?;
this.remove(position);
Ok(NoneType)
}
}
}
#[cfg(test)]
mod tests {
use crate::assert;
#[test]
fn test_error_codes() {
assert::fail(
"x = [1, 2, 3, 2]; x.remove(2); x.remove(2); x.remove(2)",
"not found in list",
);
}
#[test]
fn test_index() {
assert::fail("[True].index(True, 1, 0)", "not found");
}
#[test]
fn recursive_list() {
assert::is_true(
r#"
cyclic = [1, 2, 3]
cyclic[1] = cyclic
len(cyclic) == 3 and len(cyclic[1]) == 3
"#,
)
}
}