Add filter and zip builtins

windelbouwman · windelbouwman · commit df0400dffbe7 · 2018-11-07T09:41:16.000+01:00
diff --git a/tests/snippets/builtins.py b/tests/snippets/builtins.py
@@ -13,3 +13,7 @@
 
 assert type(frozenset) is type
 
+assert list(zip(['a', 'b', 'c'], range(3), [9, 8, 7, 99])) == [('a', 0, 9), ('b', 1, 8), ('c', 2, 7)]
+
+assert list(filter(lambda x: ((x % 2) == 0), [0, 1, 2])) == [0, 2]
+
diff --git a/tests/snippets/class.py b/tests/snippets/class.py
@@ -33,6 +33,7 @@ def foo(self):
 
 
 class Bar:
+    """ W00t """
     def __init__(self, x):
         self.x = x
 
@@ -49,6 +50,9 @@ def kungfu(x):
         assert x == 3
 
 
+# TODO:
+# assert Bar.__doc__ == " W00t "
+
 bar = Bar(42)
 
 bar.fubar(2)
diff --git a/vm/src/builtins.rs b/vm/src/builtins.rs
@@ -267,7 +267,29 @@ fn builtin_exec(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     vm.run_code_obj(code_obj, scope)
 }
 
-// builtin_filter
+fn builtin_filter(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
+    arg_check!(vm, args, required = [(function, None), (iterable, None)]);
+
+    // TODO: process one element at a time from iterators.
+    let iterable = vm.extract_elements(iterable)?;
+
+    let mut new_items = vec![];
+    for element in iterable {
+        // apply function:
+        let args = PyFuncArgs {
+            args: vec![element.clone()],
+            kwargs: vec![],
+        };
+        let result = vm.invoke(function.clone(), args)?;
+        let result = objbool::boolval(vm, result)?;
+        if result {
+            new_items.push(element);
+        }
+    }
+
+    Ok(vm.ctx.new_list(new_items))
+}
+
 // builtin_format
 
 fn builtin_getattr(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
@@ -581,7 +603,32 @@ fn builtin_sum(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
 }
 
 // builtin_vars
-// builtin_zip
+
+fn builtin_zip(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
+    no_kwargs!(vm, args);
+
+    // TODO: process one element at a time from iterators.
+    let mut iterables = vec![];
+    for iterable in args.args.iter() {
+        let iterable = vm.extract_elements(iterable)?;
+        iterables.push(iterable);
+    }
+
+    let minsize: usize = iterables.iter().map(|i| i.len()).min().unwrap_or(0);
+
+    let mut new_items = vec![];
+    for i in 0..minsize {
+        let items = iterables
+            .iter()
+            .map(|iterable| iterable[i].clone())
+            .collect();
+        let element = vm.ctx.new_tuple(items);
+        new_items.push(element);
+    }
+
+    Ok(vm.ctx.new_list(new_items))
+}
+
 // builtin___import__
 
 pub fn make_module(ctx: &PyContext) -> PyObjectRef {
@@ -616,6 +663,7 @@ pub fn make_module(ctx: &PyContext) -> PyObjectRef {
     dict.insert(String::from("exec"), ctx.new_rustfunc(builtin_exec));
     dict.insert(String::from("float"), ctx.float_type());
     dict.insert(String::from("frozenset"), ctx.frozenset_type());
+    dict.insert(String::from("filter"), ctx.new_rustfunc(builtin_filter));
     dict.insert(String::from("getattr"), ctx.new_rustfunc(builtin_getattr));
     dict.insert(String::from("hasattr"), ctx.new_rustfunc(builtin_hasattr));
     dict.insert(String::from("hash"), ctx.new_rustfunc(builtin_hash));
@@ -654,6 +702,7 @@ pub fn make_module(ctx: &PyContext) -> PyObjectRef {
     dict.insert(String::from("super"), ctx.super_type());
     dict.insert(String::from("tuple"), ctx.tuple_type());
     dict.insert(String::from("type"), ctx.type_type());
+    dict.insert(String::from("zip"), ctx.new_rustfunc(builtin_zip));
 
     // Exceptions:
     dict.insert(
diff --git a/vm/src/macros.rs b/vm/src/macros.rs
@@ -82,3 +82,15 @@ macro_rules! arg_check {
         }
     };
 }
+
+macro_rules! no_kwargs {
+    ( $vm: ident, $args:ident ) => {
+        // Zero-arg case
+        if $args.kwargs.len() != 0 {
+            return Err($vm.new_type_error(format!(
+                "Expected no keyword arguments (got: {})",
+                $args.kwargs.len()
+            )));
+        }
+    };
+}
diff --git a/vm/src/vm.rs b/vm/src/vm.rs
@@ -37,6 +37,22 @@ pub struct VirtualMachine {
 }
 
 impl VirtualMachine {
+    /// Create a new `VirtualMachine` structure.
+    pub fn new() -> VirtualMachine {
+        let ctx = PyContext::new();
+        let builtins = builtins::make_module(&ctx);
+        let sysmod = sysmodule::mk_module(&ctx);
+        // Add builtins as builtins module:
+        // sysmod.get_attr("modules").unwrap().set_item("builtins", builtins.clone());
+        let stdlib_inits = stdlib::get_module_inits();
+        VirtualMachine {
+            builtins: builtins,
+            sys_module: sysmod,
+            stdlib_inits,
+            ctx: ctx,
+        }
+    }
+
     pub fn run_code_obj(&mut self, code: PyObjectRef, scope: PyObjectRef) -> PyResult {
         let mut frame = Frame::new(code, scope);
         frame.run_frame_full(self)
@@ -122,21 +138,6 @@ impl VirtualMachine {
         &self.ctx
     }
 
-    pub fn new() -> VirtualMachine {
-        let ctx = PyContext::new();
-        let builtins = builtins::make_module(&ctx);
-        let sysmod = sysmodule::mk_module(&ctx);
-        // Add builtins as builtins module:
-        // sysmod.get_attr("modules").unwrap().set_item("builtins", builtins.clone());
-        let stdlib_inits = stdlib::get_module_inits();
-        VirtualMachine {
-            builtins: builtins,
-            sys_module: sysmod,
-            stdlib_inits,
-            ctx: ctx,
-        }
-    }
-
     pub fn get_builtin_scope(&mut self) -> PyObjectRef {
         let a2 = &*self.builtins.borrow();
         match a2.kind {
