Skip to content

Commit 91fcd77

Browse files
skinnyBatskinnyBat
committed
Make map lazy and accept multiple iterables
1 parent 8cc6821 commit 91fcd77

5 files changed

Lines changed: 66 additions & 41 deletions

File tree

‎tests/snippets/builtin_filter.py‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,3 @@
1-
21
assert list(filter(lambda x: ((x % 2) == 0), [0, 1, 2])) == [0, 2]
32

43
assert list(filter(None, [0, 1, 2])) == [0, 1, 2]

‎tests/snippets/builtin_map.py‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
a = list(map(str, [1, 2, 3]))
2+
assert a == ['1', '2', '3']
3+
4+
5+
b = list(map(lambda x, y: x + y, [1, 2, 4], [3, 5]))
6+
assert b == [4, 7]
7+
8+
9+
# test infinite iterator
10+
class Counter(object):
11+
counter = 0
12+
13+
def __next__(self):
14+
self.counter += 1
15+
return self.counter
16+
17+
def __iter__(self):
18+
return self
19+
20+
21+
it = map(lambda x: x+1, Counter())
22+
assert next(it) == 2
23+
assert next(it) == 3

‎tests/snippets/builtins.py‎

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,3 @@
1-
2-
a = list(map(str, [1, 2, 3]))
3-
assert a == ['1', '2', '3']
4-
51
x = sum(map(int, a))
62
assert x == 6
73

‎vm/src/builtins.rs‎

Lines changed: 7 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -412,30 +412,14 @@ fn builtin_locals(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
412412
}
413413

414414
fn builtin_map(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
415-
arg_check!(vm, args, required = [(function, None), (iter_target, None)]);
416-
let iterator = objiter::get_iter(vm, iter_target)?;
417-
let mut elements = vec![];
418-
loop {
419-
match vm.call_method(&iterator, "__next__", vec![]) {
420-
Ok(v) => {
421-
// Now apply function:
422-
let mapped_value = vm.invoke(
423-
function.clone(),
424-
PyFuncArgs {
425-
args: vec![v],
426-
kwargs: vec![],
427-
},
428-
)?;
429-
elements.push(mapped_value);
430-
}
431-
Err(_) => break,
432-
}
415+
no_kwargs!(vm, args);
416+
if args.args.len() < 2 {
417+
Err(vm.new_type_error("map() must have at least two arguments.".to_owned()))
418+
} else {
419+
let function = &args.args[0];
420+
let iterables = &args.args[1..];
421+
objiter::create_map(vm, function, iterables)
433422
}
434-
435-
trace!("Mapped elements: {:?}", elements);
436-
437-
// TODO: when iterators are implemented, we can improve this function.
438-
Ok(vm.ctx.new_list(elements))
439423
}
440424

441425
fn builtin_max(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {

‎vm/src/obj/objiter.rs‎

Lines changed: 36 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -84,19 +84,25 @@ pub fn create_filter(
8484
Ok(iter_obj)
8585
}
8686

87-
//pub fn create_map(vm: &mut VirtualMachine,
88-
// mapper: &PyObjectRef,
89-
// iterators: &PyObjectRef) -> PyResult {
90-
// let iter_obj = PyObject::new(
91-
// PyObjectPayload::MapIterator {
92-
// predicate: predicate.clone(),
93-
// iterator: iterator.clone(),
94-
// },
95-
// vm.ctx.iter_type(),
96-
// );
97-
//
98-
// Ok(iter_obj)
99-
//}
87+
pub fn create_map(
88+
vm: &mut VirtualMachine,
89+
mapper: &PyObjectRef,
90+
iterables: &[PyObjectRef],
91+
) -> PyResult {
92+
let iterators = iterables
93+
.into_iter()
94+
.map(|iterable| get_iter(vm, iterable))
95+
.collect::<Result<Vec<_>, _>>()?;
96+
let iter_obj = PyObject::new(
97+
PyObjectPayload::MapIterator {
98+
mapper: mapper.clone(),
99+
iterators,
100+
},
101+
vm.ctx.iter_type(),
102+
);
103+
104+
Ok(iter_obj)
105+
}
100106

101107
// Sequence iterator:
102108
fn iter_new(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
@@ -195,6 +201,23 @@ fn iter_next(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
195201
}
196202
}
197203
}
204+
PyObjectPayload::MapIterator {
205+
ref mut mapper,
206+
ref mut iterators,
207+
} => {
208+
let next_objs = iterators
209+
.iter()
210+
.map(|iterator| call_next(vm, iterator))
211+
.collect::<Result<Vec<_>, _>>()?;
212+
213+
vm.invoke(
214+
mapper.clone(),
215+
PyFuncArgs {
216+
args: next_objs,
217+
kwargs: vec![],
218+
},
219+
)
220+
}
198221
_ => {
199222
panic!("NOT IMPL");
200223
}

0 commit comments

Comments
 (0)