Make map and filter into proper types

tests/snippets/builtin_filter.py

@@ -3,6 +3,8 @@ assert list(filter(lambda x: ((x % 2) == 0), [0, 1, 2])) == [0, 2]
 # None implies identity
 assert list(filter(None, [0, 1, 2])) == [1, 2]
 
+assert type(filter(None, [])) == filter
+
 
 # test infinite iterator
 class Counter(object):

tests/snippets/builtin_map.py

@@ -5,6 +5,8 @@ assert a == ['1', '2', '3']
 b = list(map(lambda x, y: x + y, [1, 2, 4], [3, 5]))
 assert b == [4, 7]
 
+assert type(map(lambda x: x, [])) == map
+
 
 # test infinite iterator
 class Counter(object):

vm/src/builtins.rs

@@ -301,12 +301,6 @@ fn builtin_exec(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     vm.run_code_obj(code_obj, scope)
 }
 
-fn builtin_filter(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
-    arg_check!(vm, args, required = [(function, None), (iterable, None)]);
-
-    objiter::create_filter(vm, function, iterable)
-}
-
 fn builtin_format(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     arg_check!(
         vm,
@@ -411,17 +405,6 @@ fn builtin_locals(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     Ok(vm.get_locals())
 }
 
-fn builtin_map(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
-    no_kwargs!(vm, args);
-    if args.args.len() < 2 {
-        Err(vm.new_type_error("map() must have at least two arguments.".to_owned()))
-    } else {
-        let function = &args.args[0];
-        let iterables = &args.args[1..];
-        objiter::create_map(vm, function, iterables)
-    }
-}
-
 fn builtin_max(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     let candidates = if args.args.len() > 1 {
         args.args.clone()
@@ -716,7 +699,7 @@ pub fn make_module(ctx: &PyContext) -> PyObjectRef {
     ctx.set_attr(&py_mod, "exec", ctx.new_rustfunc(builtin_exec));
     ctx.set_attr(&py_mod, "float", ctx.float_type());
     ctx.set_attr(&py_mod, "frozenset", ctx.frozenset_type());
-    ctx.set_attr(&py_mod, "filter", ctx.new_rustfunc(builtin_filter));
+    ctx.set_attr(&py_mod, "filter", ctx.filter_type());
     ctx.set_attr(&py_mod, "format", ctx.new_rustfunc(builtin_format));
     ctx.set_attr(&py_mod, "getattr", ctx.new_rustfunc(builtin_getattr));
     ctx.set_attr(&py_mod, "hasattr", ctx.new_rustfunc(builtin_hasattr));
@@ -730,7 +713,7 @@ pub fn make_module(ctx: &PyContext) -> PyObjectRef {
     ctx.set_attr(&py_mod, "len", ctx.new_rustfunc(builtin_len));
     ctx.set_attr(&py_mod, "list", ctx.list_type());
     ctx.set_attr(&py_mod, "locals", ctx.new_rustfunc(builtin_locals));
-    ctx.set_attr(&py_mod, "map", ctx.new_rustfunc(builtin_map));
+    ctx.set_attr(&py_mod, "map", ctx.map_type());
     ctx.set_attr(&py_mod, "max", ctx.new_rustfunc(builtin_max));
     ctx.set_attr(&py_mod, "memoryview", ctx.memoryview_type());
     ctx.set_attr(&py_mod, "min", ctx.new_rustfunc(builtin_min));

vm/src/obj/mod.rs

@@ -6,6 +6,7 @@ pub mod objbytes;
 pub mod objcode;
 pub mod objcomplex;
 pub mod objdict;
+pub mod objfilter;
 pub mod objfloat;
 pub mod objframe;
 pub mod objfunction;
@@ -13,6 +14,7 @@ pub mod objgenerator;
 pub mod objint;
 pub mod objiter;
 pub mod objlist;
+pub mod objmap;
 pub mod objmemory;
 pub mod objobject;
 pub mod objproperty;

vm/src/obj/objfilter.rs

@@ -0,0 +1,83 @@
+use super::super::pyobject::{
+    IdProtocol, PyContext, PyFuncArgs, PyObject, PyObjectPayload, PyObjectRef, PyResult,
+    TypeProtocol,
+};
+use super::super::vm::VirtualMachine;
+use super::objbool;
+use super::objiter;
+use super::objtype; // Required for arg_check! to use isinstance
+
+pub fn filter_new(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
+    arg_check!(
+        vm,
+        args,
+        required = [(cls, None), (function, None), (iterable, None)]
+    );
+    let iterator = objiter::get_iter(vm, iterable)?;
+    Ok(PyObject::new(
+        PyObjectPayload::FilterIterator {
+            predicate: function.clone(),
+            iterator,
+        },
+        cls.clone(),
+    ))
+}
+
+fn filter_iter(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
+    arg_check!(vm, args, required = [(filter, Some(vm.ctx.filter_type()))]);
+    // Return self:
+    Ok(filter.clone())
+}
+
+fn filter_contains(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
+    arg_check!(
+        vm,
+        args,
+        required = [(filter, Some(vm.ctx.filter_type())), (needle, None)]
+    );
+    objiter::contains(vm, filter, needle)
+}
+
+fn filter_next(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
+    arg_check!(vm, args, required = [(filter, Some(vm.ctx.filter_type()))]);
+
+    if let PyObjectPayload::FilterIterator {
+        ref mut predicate,
+        ref mut iterator,
+    } = filter.borrow_mut().payload
+    {
+        loop {
+            let next_obj = objiter::call_next(vm, iterator)?;
+            let predicate_value = if predicate.is(&vm.get_none()) {
+                next_obj.clone()
+            } else {
+                // the predicate itself can raise StopIteration which does stop the filter
+                // iteration
+                vm.invoke(
+                    predicate.clone(),
+                    PyFuncArgs {
+                        args: vec![next_obj.clone()],
+                        kwargs: vec![],
+                    },
+                )?
+            };
+            if objbool::boolval(vm, predicate_value)? {
+                return Ok(next_obj);
+            }
+        }
+    } else {
+        panic!("filter doesn't have correct payload");
+    }
+}
+
+pub fn init(context: &PyContext) {
+    let filter_type = &context.filter_type;
+    context.set_attr(
+        &filter_type,
+        "__contains__",
+        context.new_rustfunc(filter_contains),
+    );
+    context.set_attr(&filter_type, "__iter__", context.new_rustfunc(filter_iter));
+    context.set_attr(&filter_type, "__new__", context.new_rustfunc(filter_new));
+    context.set_attr(&filter_type, "__next__", context.new_rustfunc(filter_next));
+}

vm/src/obj/objiter.rs

@@ -3,8 +3,7 @@
  */
 
 use super::super::pyobject::{
-    IdProtocol, PyContext, PyFuncArgs, PyObject, PyObjectPayload, PyObjectRef, PyResult,
-    TypeProtocol,
+    PyContext, PyFuncArgs, PyObjectPayload, PyObjectRef, PyResult, TypeProtocol,
 };
 use super::super::vm::VirtualMachine;
 use super::objbool;
@@ -66,42 +65,19 @@ pub fn get_all(
     Ok(elements)
 }
 
-// Should filter/map have their own class?
-pub fn create_filter(
-    vm: &mut VirtualMachine,
-    predicate: &PyObjectRef,
-    iterable: &PyObjectRef,
-) -> PyResult {
-    let iterator = get_iter(vm, iterable)?;
-    let iter_obj = PyObject::new(
-        PyObjectPayload::FilterIterator {
-            predicate: predicate.clone(),
-            iterator,
-        },
-        vm.ctx.iter_type(),
-    );
-
-    Ok(iter_obj)
-}
-
-pub fn create_map(
-    vm: &mut VirtualMachine,
-    mapper: &PyObjectRef,
-    iterables: &[PyObjectRef],
-) -> PyResult {
-    let iterators = iterables
-        .into_iter()
-        .map(|iterable| get_iter(vm, iterable))
-        .collect::<Result<Vec<_>, _>>()?;
-    let iter_obj = PyObject::new(
-        PyObjectPayload::MapIterator {
-            mapper: mapper.clone(),
-            iterators,
-        },
-        vm.ctx.iter_type(),
-    );
-
-    Ok(iter_obj)
+pub fn contains(vm: &mut VirtualMachine, iter: &PyObjectRef, needle: &PyObjectRef) -> PyResult {
+    loop {
+        if let Some(element) = get_next_object(vm, iter)? {
+            let equal = vm.call_method(needle, "__eq__", vec![element.clone()])?;
+            if objbool::get_value(&equal) {
+                return Ok(vm.new_bool(true));
+            } else {
+                continue;
+            }
+        } else {
+            return Ok(vm.new_bool(false));
+        }
+    }
 }
 
 // Sequence iterator:
@@ -123,108 +99,49 @@ fn iter_contains(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
         args,
         required = [(iter, Some(vm.ctx.iter_type())), (needle, None)]
     );
-    loop {
-        match vm.call_method(&iter, "__next__", vec![]) {
-            Ok(element) => match vm.call_method(needle, "__eq__", vec![element.clone()]) {
-                Ok(value) => {
-                    if objbool::get_value(&value) {
-                        return Ok(vm.new_bool(true));
-                    } else {
-                        continue;
-                    }
-                }
-                Err(_) => return Err(vm.new_type_error("".to_string())),
-            },
-            Err(_) => return Ok(vm.new_bool(false)),
-        }
-    }
+    contains(vm, iter, needle)
 }
 
 fn iter_next(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     arg_check!(vm, args, required = [(iter, Some(vm.ctx.iter_type()))]);
 
-    let ref mut payload = iter.borrow_mut().payload;
-    match payload {
-        PyObjectPayload::Iterator {
-            ref mut position,
-            iterated_obj: ref mut iterated_obj_ref,
-        } => {
-            let iterated_obj = iterated_obj_ref.borrow_mut();
-            match iterated_obj.payload {
-                PyObjectPayload::Sequence { ref elements } => {
-                    if *position < elements.len() {
-                        let obj_ref = elements[*position].clone();
-                        *position += 1;
-                        Ok(obj_ref)
-                    } else {
-                        let stop_iteration_type = vm.ctx.exceptions.stop_iteration.clone();
-                        let stop_iteration =
-                            vm.new_exception(stop_iteration_type, "End of iterator".to_string());
-                        Err(stop_iteration)
-                    }
-                }
-
-                PyObjectPayload::Range { ref range } => {
-                    if let Some(int) = range.get(BigInt::from(*position)) {
-                        *position += 1;
-                        Ok(vm.ctx.new_int(int.to_bigint().unwrap()))
-                    } else {
-                        let stop_iteration_type = vm.ctx.exceptions.stop_iteration.clone();
-                        let stop_iteration =
-                            vm.new_exception(stop_iteration_type, "End of iterator".to_string());
-                        Err(stop_iteration)
-                    }
-                }
-                _ => {
-                    panic!("NOT IMPL");
-                }
-            }
-        }
-        PyObjectPayload::FilterIterator {
-            ref mut predicate,
-            ref mut iterator,
-        } => {
-            loop {
-                let next_obj = call_next(vm, iterator)?;
-                let predicate_value = if predicate.is(&vm.get_none()) {
-                    next_obj.clone()
+    if let PyObjectPayload::Iterator {
+        ref mut position,
+        iterated_obj: ref mut iterated_obj_ref,
+    } = iter.borrow_mut().payload
+    {
+        let iterated_obj = iterated_obj_ref.borrow_mut();
+        match iterated_obj.payload {
+            PyObjectPayload::Sequence { ref elements } => {
+                if *position < elements.len() {
+                    let obj_ref = elements[*position].clone();
+                    *position += 1;
+                    Ok(obj_ref)
                 } else {
-                    // the predicate itself can raise StopIteration which does stop the filter
-                    // iteration
-                    vm.invoke(
-                        predicate.clone(),
-                        PyFuncArgs {
-                            args: vec![next_obj.clone()],
-                            kwargs: vec![],
-                        },
-                    )?
-                };
-                if objbool::boolval(vm, predicate_value)? {
-                    return Ok(next_obj);
+                    let stop_iteration_type = vm.ctx.exceptions.stop_iteration.clone();
+                    let stop_iteration =
+                        vm.new_exception(stop_iteration_type, "End of iterator".to_string());
+                    Err(stop_iteration)
                 }
             }
-        }
-        PyObjectPayload::MapIterator {
-            ref mut mapper,
-            ref mut iterators,
-        } => {
-            let next_objs = iterators
-                .iter()
-                .map(|iterator| call_next(vm, iterator))
-                .collect::<Result<Vec<_>, _>>()?;
 
-            // the mapper itself can raise StopIteration which does stop the map iteration
-            vm.invoke(
-                mapper.clone(),
-                PyFuncArgs {
-                    args: next_objs,
-                    kwargs: vec![],
-                },
-            )
-        }
-        _ => {
-            panic!("NOT IMPL");
+            PyObjectPayload::Range { ref range } => {
+                if let Some(int) = range.get(BigInt::from(*position)) {
+                    *position += 1;
+                    Ok(vm.ctx.new_int(int.to_bigint().unwrap()))
+                } else {
+                    let stop_iteration_type = vm.ctx.exceptions.stop_iteration.clone();
+                    let stop_iteration =
+                        vm.new_exception(stop_iteration_type, "End of iterator".to_string());
+                    Err(stop_iteration)
+                }
+            }
+            _ => {
+                panic!("NOT IMPL");
+            }
         }
+    } else {
+        panic!("NOT IMPL");
     }
 }
 

vm/src/obj/objmap.rs

@@ -0,0 +1,81 @@
+use super::super::pyobject::{
+    PyContext, PyFuncArgs, PyObject, PyObjectPayload, PyObjectRef, PyResult, TypeProtocol,
+};
+use super::super::vm::VirtualMachine;
+use super::objiter;
+use super::objtype; // Required for arg_check! to use isinstance
+
+pub fn map_new(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
+    no_kwargs!(vm, args);
+    let cls = &args.args[0];
+    if args.args.len() < 3 {
+        Err(vm.new_type_error("map() must have at least two arguments.".to_owned()))
+    } else {
+        let function = &args.args[1];
+        let iterables = &args.args[2..];
+        let iterators = iterables
+            .into_iter()
+            .map(|iterable| objiter::get_iter(vm, iterable))
+            .collect::<Result<Vec<_>, _>>()?;
+        Ok(PyObject::new(
+            PyObjectPayload::MapIterator {
+                mapper: function.clone(),
+                iterators,
+            },
+            cls.clone(),
+        ))
+    }
+}
+
+fn map_iter(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
+    arg_check!(vm, args, required = [(map, Some(vm.ctx.map_type()))]);
+    // Return self:
+    Ok(map.clone())
+}
+
+fn map_contains(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
+    arg_check!(
+        vm,
+        args,
+        required = [(map, Some(vm.ctx.map_type())), (needle, None)]
+    );
+    objiter::contains(vm, map, needle)
+}
+
+fn map_next(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
+    arg_check!(vm, args, required = [(map, Some(vm.ctx.map_type()))]);
+
+    if let PyObjectPayload::MapIterator {
+        ref mut mapper,
+        ref mut iterators,
+    } = map.borrow_mut().payload
+    {
+        let next_objs = iterators
+            .iter()
+            .map(|iterator| objiter::call_next(vm, iterator))
+            .collect::<Result<Vec<_>, _>>()?;
+
+        // the mapper itself can raise StopIteration which does stop the map iteration
+        vm.invoke(
+            mapper.clone(),
+            PyFuncArgs {
+                args: next_objs,
+                kwargs: vec![],
+            },
+        )
+    } else {
+        panic!("map doesn't have correct payload");
+    }
+}
+
+pub fn init(context: &PyContext) {
+    let map_type = &context.map_type;
+    context.set_attr(
+        &map_type,
+        "__contains__",
+        context.new_rustfunc(map_contains),
+    );
+    context.set_attr(&map_type, "__iter__", context.new_rustfunc(map_iter));
+    context.set_attr(&map_type, "__new__", context.new_rustfunc(map_new));
+    context.set_attr(&map_type, "__next__", context.new_rustfunc(map_next));
+}

vm/src/pyobject.rs

@@ -7,6 +7,7 @@ use super::obj::objbytes;
 use super::obj::objcode;
 use super::obj::objcomplex;
 use super::obj::objdict;
+use super::obj::objfilter;
 use super::obj::objfloat;
 use super::obj::objframe;
 use super::obj::objfunction;
@@ -14,6 +15,7 @@ use super::obj::objgenerator;
 use super::obj::objint;
 use super::obj::objiter;
 use super::obj::objlist;
+use super::obj::objmap;
 use super::obj::objmemory;
 use super::obj::objobject;
 use super::obj::objproperty;
@@ -106,6 +108,7 @@ pub struct PyContext {
     pub classmethod_type: PyObjectRef,
     pub code_type: PyObjectRef,
     pub dict_type: PyObjectRef,
+    pub filter_type: PyObjectRef,
     pub float_type: PyObjectRef,
     pub frame_type: PyObjectRef,
     pub frozenset_type: PyObjectRef,
@@ -116,6 +119,7 @@ pub struct PyContext {
     pub true_value: PyObjectRef,
     pub false_value: PyObjectRef,
     pub list_type: PyObjectRef,
+    pub map_type: PyObjectRef,
     pub memoryview_type: PyObjectRef,
     pub none: PyObjectRef,
     pub tuple_type: PyObjectRef,
@@ -200,6 +204,8 @@ impl PyContext {
         let bytearray_type = create_type("bytearray", &type_type, &object_type, &dict_type);
         let tuple_type = create_type("tuple", &type_type, &object_type, &dict_type);
         let iter_type = create_type("iter", &type_type, &object_type, &dict_type);
+        let filter_type = create_type("filter", &type_type, &object_type, &dict_type);
+        let map_type = create_type("map", &type_type, &object_type, &dict_type);
         let bool_type = create_type("bool", &type_type, &int_type, &dict_type);
         let memoryview_type = create_type("memoryview", &type_type, &object_type, &dict_type);
         let code_type = create_type("code", &type_type, &int_type, &dict_type);
@@ -240,6 +246,8 @@ impl PyContext {
             false_value,
             tuple_type,
             iter_type,
+            filter_type,
+            map_type,
             dict_type,
             none: none,
             str_type: str_type,
@@ -275,6 +283,8 @@ impl PyContext {
         objsuper::init(&context);
         objtuple::init(&context);
         objiter::init(&context);
+        objfilter::init(&context);
+        objmap::init(&context);
         objbool::init(&context);
         objcode::init(&context);
         objframe::init(&context);
@@ -346,6 +356,14 @@ impl PyContext {
         self.iter_type.clone()
     }
 
+    pub fn filter_type(&self) -> PyObjectRef {
+        self.filter_type.clone()
+    }
+
+    pub fn map_type(&self) -> PyObjectRef {
+        self.map_type.clone()
+    }
+
     pub fn str_type(&self) -> PyObjectRef {
         self.str_type.clone()
     }