Refactor map function in Dict to improve parameter handling and update related tests for consistency

This commit is contained in:
Wesley Yang 2025-04-03 18:18:11 +08:00
parent c9c921b877
commit 06bc7bac66
2 changed files with 35 additions and 10 deletions

View File

@ -264,8 +264,8 @@ impl Dict {
///
/// ```example
/// #let prices = (apples: 2, oranges: 3, bananas: 1.5)
/// #prices.map(key => key.len()) \
/// #prices.map((key, price) => (key, price * 1.1))
/// #prices.map(pair => pair.at(0).len())
/// #prices.map((key, value) => (key, value * 1.1))
/// ```
#[func]
pub fn map(
@ -275,6 +275,7 @@ impl Dict {
/// The function to apply to each key-value pair.
/// The function can either take a single parameter (receiving a pair as array of length 2),
/// or two parameters (receiving key and value separately).
/// Parameters exceeding two will be ignored.
mapper: Func,
) -> SourceResult<Value> {
let mut dict_result = IndexMap::new();
@ -282,18 +283,42 @@ impl Dict {
let mut is_dict = true;
// try to check the number of parameters, if not, use array form
let use_two_args = mapper.params().is_some_and(|params| params.len() >= 2);
let mut first_pair = true;
let mut use_single_arg = false;
for (key, value) in self {
// choose how to pass parameters based on the function signature
let mapped = if use_two_args {
mapper.call(engine, context, [Value::Str(key.clone()), value.clone()])?
} else {
let mapped = if first_pair {
// try two calling ways for the first pair
first_pair = false;
// try to call with two parameters
let result = mapper.call(
engine,
context,
[Value::Str(key.clone()), value.clone()],
);
// if failed, try to call with one parameter
if result.is_err() {
use_single_arg = true;
mapper.call(
engine,
context,
[Value::Array(array![Value::Str(key.clone()), value])],
)?
} else {
result?
}
} else if use_single_arg {
// try to call with one parameter
mapper.call(
engine,
context,
[Value::Array(array![Value::Str(key.clone()), value])],
)?
} else {
// try to call with two parameters
mapper.call(engine, context, [Value::Str(key.clone()), value.clone()])?
};
// check if the result is a dictionary key-value pair

View File

@ -29,13 +29,13 @@
// test map return new dict
#test(
dict.map(((key, value)) => (key, value * 2)),
dict.map((key, value) => (key, value * 2)),
(a: 2, b: 4, c: 6)
)
// test map empty dict
#test(
(:).map(((key, value)) => (key, value * 2)),
(:).map((key, value) => (key, value * 2)),
(:)
)
@ -47,7 +47,7 @@
// test map return array(different return type)
#test(
dict.map(((key, value)) => if value > 1 { (key, value * 2) } else { "key smaller than 1: " + key }),
dict.map((key, value) => if value > 1 { (key, value * 2) } else { "key smaller than 1: " + key }),
("key smaller than 1: a", ("b", 4), ("c", 6))
)