Add support for keyword arguments and **kwargs.

progval · progval · commit e380b9c73aab · 2016-05-08T13:42:32.000+02:00
diff --git a/functional_tests/arguments.py b/functional_tests/arguments.py
@@ -0,0 +1,17 @@
+def f(foo, bar, *baz, qux, quux, fiz='buz', **kwargs):
+    print('foo:', foo)
+    print('bar:', bar)
+    for x in baz:
+        print('baz:', x)
+    print('qux:', qux)
+    print('quux:', quux)
+    print('fiz:', fiz)
+    # TODO: print kwargs
+    print('--')
+
+f(1, 2, qux=3, quux=4)
+f(1, 2, 3, qux=4, quux=5)
+f(1, 2, 3, 21, qux=4, quux=5)
+f(1, 2, qux=3, quux=4, fiz=5)
+f(1, 2, 3, qux=4, quux=5, fiz=6)
+f(1, 2, 3, 21, qux=4, quux=5, fiz=6)
diff --git a/src/marshal/decode.rs b/src/marshal/decode.rs
@@ -224,7 +224,7 @@ pub fn read_object<R: io::Read>(reader: &mut R, store: &mut ObjectStore, primiti
             let lnotab = try!(read_object(reader, store, primitive_objects, references)); // TODO: decode this
             let code = Code {
                 argcount: argcount as usize,
-                kwonlyargcount: kwonlyargcount,
+                kwonlyargcount: kwonlyargcount as usize,
                 nlocals: nlocals,
                 stacksize: stacksize,
                 flags: flags,
diff --git a/src/objects/mod.rs b/src/objects/mod.rs
@@ -1,5 +1,6 @@
 extern crate itertools;
 use std::collections::HashMap;
+use std::collections::HashSet;
 use std::sync::atomic::{AtomicUsize, Ordering, ATOMIC_USIZE_INIT};
 use std::fmt;
 use self::itertools::Itertools;
@@ -12,7 +13,7 @@ use super::sandbox::EnvProxy;
 #[derive(Eq)]
 pub struct Code {
     pub argcount: usize,
-    pub kwonlyargcount: u32,
+    pub kwonlyargcount: usize,
     pub nlocals: u32,
     pub stacksize: u32,
     pub flags: u32,
@@ -32,6 +33,36 @@ impl Code {
     pub fn co_varargs(&self) -> bool {
         self.flags & 0x4 != 0
     }
+    pub fn co_varkwargs(&self) -> bool {
+        self.flags & 0x8 != 0
+    }
+    pub fn get_varargs_name(&self) -> Option<&String> {
+        if self.co_varargs() {
+            Some(self.varnames.get(self.argcount+self.kwonlyargcount).unwrap())
+        }
+        else {
+            None
+        }
+    }
+    pub fn get_varkwargs_name(&self) -> Option<&String> {
+        if self.co_varkwargs() {
+            let mut index = self.argcount+self.kwonlyargcount;
+            if self.co_varkwargs() {
+                index += 1;
+            }
+            Some(self.varnames.get(index).unwrap())
+        }
+        else {
+            None
+        }
+    }
+    pub fn keywords(&self) -> HashSet<&String> {
+        let mut res = HashSet::new();
+        for name in self.varnames[self.argcount..(self.argcount+self.kwonlyargcount)].iter() {
+            res.insert(name);
+        }
+        res
+    }
 }
 
 #[derive(Debug)]
@@ -48,9 +79,10 @@ pub enum ObjectContent {
     List(Vec<ObjectRef>),
     Code(Box<Code>),
     Set(Vec<ObjectRef>),
+    Dict(Vec<(ObjectRef, ObjectRef)>),
     FrozenSet(Vec<ObjectRef>),
     Bytes(Vec<u8>),
-    Function(String, ObjectRef), // module, code
+    Function(String, ObjectRef, HashMap<String, ObjectRef>), // module, code, default arguments
     Module(ObjectRef),
     PrimitiveNamespace, // __primitives__
     PrimitiveFunction(String),
@@ -131,12 +163,13 @@ impl ObjectRef {
             ObjectContent::Int(ref i) => i.to_string(),
             ObjectContent::Bytes(ref s) => "<bytes>".to_string(), // TODO
             ObjectContent::String(ref s) => format!("'{}'", s), // TODO: escape
+            ObjectContent::Dict(ref l) => format!("{{{}}}", l.iter().map(|arg| { let (ref k, ref v) = *arg; format!("{}: {}", k.repr(store), v.repr(store))}).join(", ")),
             ObjectContent::Tuple(ref l) => format!("tuple({})", ObjectRef::repr_vec(l, store)),
             ObjectContent::List(ref l) => format!("[{}]", ObjectRef::repr_vec(l, store)),
             ObjectContent::Code(_) => "<code object>".to_string(),
             ObjectContent::Set(ref l) => format!("set({})", ObjectRef::repr_vec(l, store)),
             ObjectContent::FrozenSet(ref l) => format!("frozenset({})", ObjectRef::repr_vec(l, store)),
-            ObjectContent::Function(ref module, ref _code) => {
+            ObjectContent::Function(ref module, ref _code, ref _defaults) => {
                 match obj.name {
                     None => format!("<anonymous function in module {}>", module),
                     Some(ref s) => format!("<function {} in module {}>", s, module),
@@ -167,7 +200,7 @@ impl ObjectRef {
         let func = store.deref(self);
         let ref name = func.name;
         match func.content {
-            ObjectContent::Function(ref module_name, ref _code) => module_name.clone(),
+            ObjectContent::Function(ref module_name, ref _code, ref _defaults) => module_name.clone(),
             ObjectContent::Module(ref _code) => name.clone().unwrap(),
             _ => panic!(format!("Not a function/module: {:?}", func)),
         }
@@ -241,6 +274,7 @@ pub struct PrimitiveObjects {
 
     pub set_type: ObjectRef,
     pub frozenset_type: ObjectRef,
+    pub dict_type: ObjectRef,
 
     pub bytes_type: ObjectRef,
     pub str_type: ObjectRef,
@@ -300,6 +334,7 @@ impl PrimitiveObjects {
         let list_type = store.allocate(Object::new_class("list".to_string(), None, type_ref.clone(), vec![obj_ref.clone()]));
         let set_type = store.allocate(Object::new_class("set".to_string(), None, type_ref.clone(), vec![obj_ref.clone()]));
         let frozenset_type = store.allocate(Object::new_class("frozenset".to_string(), None, type_ref.clone(), vec![obj_ref.clone()]));
+        let dict_type = store.allocate(Object::new_class("dict".to_string(), None, type_ref.clone(), vec![obj_ref.clone()]));
         let bytes_type = store.allocate(Object::new_class("bytes".to_string(), None, type_ref.clone(), vec![obj_ref.clone()]));
         let str_type = store.allocate(Object::new_class("str".to_string(), None, type_ref.clone(), vec![obj_ref.clone()]));
         let iterator_type = store.allocate(Object::new_class("iterator".to_string(), None, type_ref.clone(), vec![obj_ref.clone()]));
@@ -335,6 +370,7 @@ impl PrimitiveObjects {
         map.insert("list".to_string(), list_type.clone());
         map.insert("set".to_string(), set_type.clone());
         map.insert("frozenset".to_string(), frozenset_type.clone());
+        map.insert("dict".to_string(), dict_type.clone());
         map.insert("bytes".to_string(), bytes_type.clone());
         map.insert("str".to_string(), str_type.clone());
         map.insert("function".to_string(), function_type.clone());
@@ -359,7 +395,7 @@ impl PrimitiveObjects {
             none_type: none_type, none: none,
             int_type: int_type, bool_type: bool_type, true_obj: true_obj, false_obj: false_obj,
             tuple_type: tuple_type, list_type: list_type,
-            set_type: set_type, frozenset_type: frozenset_type,
+            set_type: set_type, frozenset_type: frozenset_type, dict_type: dict_type,
             bytes_type: bytes_type, str_type: str_type,
             iterator_type: iterator_type,
             function_type: function_type, code_type: code_type,
@@ -389,14 +425,17 @@ impl PrimitiveObjects {
     pub fn new_set(&self, v: Vec<ObjectRef>) -> Object {
         Object::new_instance(None, self.set_type.clone(), ObjectContent::Set(v))
     }
+    pub fn new_dict(&self, v: Vec<(ObjectRef, ObjectRef)>) -> Object {
+        Object::new_instance(None, self.dict_type.clone(), ObjectContent::Dict(v))
+    }
     pub fn new_frozenset(&self, v: Vec<ObjectRef>) -> Object {
         Object::new_instance(None, self.frozenset_type.clone(), ObjectContent::FrozenSet(v))
     }
     pub fn new_code(&self, c: Code) -> Object {
         Object::new_instance(None, self.code_type.clone(), ObjectContent::Code(Box::new(c)))
     }
-    pub fn new_function(&self, name: String, module_name: String, code: ObjectRef) -> Object {
-        Object::new_instance(Some(name), self.function_type.clone(), ObjectContent::Function(module_name, code))
+    pub fn new_function(&self, name: String, module_name: String, code: ObjectRef, defaults: HashMap<String, ObjectRef>) -> Object {
+        Object::new_instance(Some(name), self.function_type.clone(), ObjectContent::Function(module_name, code, defaults))
     }
     pub fn new_module(&self, name: String, code: ObjectRef) -> Object {
         Object::new_instance(Some(name), self.module.clone(), ObjectContent::Module(code))
diff --git a/src/primitives/mod.rs b/src/primitives/mod.rs
@@ -58,7 +58,7 @@ fn build_class<EP: EnvProxy>(processor: &mut State<EP>, call_stack: &mut Vec<Fra
     let mut args_iter = args.into_iter();
     parse_first_arguments!("__primitives__.build_class", processor.store, args, args_iter,
         "func" "a function": {
-            ObjectContent::Function(_, ref code_arg) => {
+            ObjectContent::Function(_, ref code_arg, _) => {
                 code = code_arg.clone();
             },
         },
diff --git a/src/processor/mod.rs b/src/processor/mod.rs
@@ -10,6 +10,7 @@ use std::collections::HashMap;
 use std::io::Read;
 use std::cell::RefCell;
 use std::rc::Rc;
+use std::iter::FromIterator;
 use super::marshal;
 use super::state::{State, PyResult, unwind, raise, return_value};
 use super::sandbox::EnvProxy;
@@ -129,7 +130,7 @@ fn load_attr<EP: EnvProxy>(state: &mut State<EP>, obj: &Object, name: &String) -
 }
 
 // Call a primitive / function / code object, with arguments.
-fn call_function<EP: EnvProxy>(state: &mut State<EP>, call_stack: &mut Vec<Frame>, func_ref: &ObjectRef, mut args: Vec<ObjectRef>, kwargs: Vec<ObjectRef>) {
+fn call_function<EP: EnvProxy>(state: &mut State<EP>, call_stack: &mut Vec<Frame>, func_ref: &ObjectRef, mut args: Vec<ObjectRef>, kwargs: Vec<(ObjectRef, ObjectRef)>) {
     // TODO: clone only if necessary
     match state.store.deref(func_ref).content.clone() {
         ObjectContent::Class(None) => {
@@ -141,33 +142,69 @@ fn call_function<EP: EnvProxy>(state: &mut State<EP>, call_stack: &mut Vec<Frame
             let frame = call_stack.last_mut().unwrap();
             frame.var_stack.push(state.store.allocate(Object::new_instance(None, func_ref.clone(), ObjectContent::OtherObject)))
         },
-        ObjectContent::Function(ref func_module, ref code_ref) => {
+        ObjectContent::Function(ref func_module, ref code_ref, ref defaults) => {
             let code = state.store.deref(code_ref).content.clone();
             if let ObjectContent::Code(code) = code {
-                if code.co_varargs() { // If it has a *args argument
+
+                let mut locals = Rc::new(RefCell::new(defaults.clone()));
+
+                if let Some(starargs_name) = code.get_varargs_name() { // If it has a *args argument
                     if code.argcount > args.len() {
                         panic!(format!("{}() takes at least {} arguments, but {} was/were given.", code.name, code.argcount, args.len()))
                     };
                     let to_vararg = args.drain(code.argcount..).collect();
                     let obj_ref = state.store.allocate(state.primitive_objects.new_tuple(to_vararg));
-                    args.push(obj_ref);
+
+                    // Bind *args
+                    assert_eq!(None, locals.borrow_mut().insert(starargs_name.clone(), obj_ref));
                 }
-                else if code.argcount != args.len() {
+                else if code.argcount != args.len() { // If it has no *args argument
                     panic!(format!("{}() takes {} arguments, but {} was/were given.", code.name, code.argcount, args.len()))
                 };
-                let mut locals = Rc::new(RefCell::new(HashMap::new()));
+
+                // Handle keyword arguments
+                let mut remaining_kwargs = vec![]; // arguments that will go to **kwargs
+                {
+                    let explicit_keywords = code.keywords();
+                    let mut locals = locals.borrow_mut();
+                    for (key, value) in kwargs.into_iter() {
+                        let key_str = match state.store.deref(&key).content {
+                            ObjectContent::String(ref s) => s,
+                            _ => panic!("Keyword names should be strings."),
+                        };
+                        if explicit_keywords.contains(key_str) {
+                            locals.insert(key_str.clone(), value);
+                        }
+                        else {
+                            remaining_kwargs.push((key, value))
+                        }
+                    }
+                }
+
+                if let Some(starkwargs_name) = code.get_varkwargs_name() { // If it has a **kwargs argument
+                    let obj_ref = state.store.allocate(state.primitive_objects.new_dict(remaining_kwargs));
+                    locals.borrow_mut().insert(starkwargs_name.clone(), obj_ref);
+                }
+                else { // If it has no **kwargs argument
+                    if remaining_kwargs.len() != 0 {
+                        panic!(format!("Unknown keyword argument to function {} with no **kwargs", code.name))
+                    }
+                }
+
+                // Bind positional arguments
                 {
                     let mut locals = locals.borrow_mut();
                     for (argname, argvalue) in code.varnames.iter().zip(args) {
                         locals.insert(argname.clone(), argvalue);
                     };
                 }
+
                 let new_frame = Frame::new(func_ref.clone(), *code, locals);
                 call_stack.push(new_frame);
             }
             else {
                 let exc = state.primitive_objects.processorerror.clone();
-                let repr = func_ref.repr(&state.store);
+                let repr = code_ref.repr(&state.store);
                 raise(state, call_stack, exc, format!("Not a code object {}", repr));
             }
         },
@@ -460,13 +497,13 @@ fn run_code<EP: EnvProxy>(state: &mut State<EP>, call_stack: &mut Vec<Frame>) ->
                 let mut func;
                 {
                     let frame = call_stack.last_mut().unwrap();
-                    kwargs = py_unwrap!(state, frame.var_stack.pop_many(nb_kwargs*2), ProcessorError::StackTooSmall);
+                    kwargs = py_unwrap!(state, frame.var_stack.pop_n_pairs(nb_kwargs), ProcessorError::StackTooSmall);
                     args = py_unwrap!(state, frame.var_stack.pop_many(nb_args), ProcessorError::StackTooSmall);
                     func = pop_stack!(state, frame.var_stack);
                 }
                 call_function(state, call_stack, &func, args, kwargs)
             },
-            Instruction::MakeFunction(0, 0, 0) => {
+            Instruction::MakeFunction(0, nb_default_kwargs, 0) => {
                 // TODO: consume default arguments and annotations
                 let obj = {
                     let frame = call_stack.last_mut().unwrap();
@@ -483,7 +520,16 @@ fn run_code<EP: EnvProxy>(state: &mut State<EP>, call_stack: &mut Vec<Frame>) ->
                 };
                 let frame = call_stack.last_mut().unwrap();
                 let code = pop_stack!(state, frame.var_stack);
-                let func = state.primitive_objects.new_function(func_name, frame.object.module(&state.store), code);
+                let raw_kwdefaults = py_unwrap!(state, frame.var_stack.pop_n_pairs(nb_default_kwargs), ProcessorError::StackTooSmall);
+                let mut kwdefaults: HashMap<String, ObjectRef> = HashMap::new();
+                kwdefaults.reserve(nb_default_kwargs);
+                for (key, value) in raw_kwdefaults {
+                    match state.store.deref(&key).content {
+                        ObjectContent::String(ref s) => { kwdefaults.insert(s.clone(), value); },
+                        _ => panic!("Defaults' keys must be strings."),
+                    }
+                }
+                let func = state.primitive_objects.new_function(func_name, frame.object.module(&state.store), code, kwdefaults);
                 frame.var_stack.push(state.store.allocate(func))
             },
             _ => panic!(format!("todo: instruction {:?}", instruction)),
diff --git a/src/varstack.rs b/src/varstack.rs
@@ -8,6 +8,7 @@ pub trait VarStack : Debug {
     fn pop_many(&mut self, count: usize) -> Option<Vec<Self::Item>>;
     fn push(&mut self, value: Self::Item);
     fn pop_all_and_get_n_last(&mut self, nb: usize) -> Option<Vec<Self::Item>>;
+    fn pop_n_pairs(&mut self, nb: usize) -> Option<Vec<(Self::Item, Self::Item)>>;
 }
 
 #[derive(Debug)]
@@ -27,7 +28,7 @@ impl<Item> VectorVarStack<Item> {
     }
 }
 
-impl<Item> VarStack for VectorVarStack<Item> where Item: Debug {
+impl<Item: Clone> VarStack for VectorVarStack<Item> where Item: Debug {
     type Item = Item;
 
     fn top(&self) -> Option<&Self::Item> {
@@ -61,4 +62,17 @@ impl<Item> VarStack for VectorVarStack<Item> where Item: Debug {
             Some(self.vector.drain(..).collect())
         }
     }
+
+    fn pop_n_pairs(&mut self, nb: usize) -> Option<Vec<(Self::Item, Self::Item)>> {
+        self.pop_many(nb*2).map(|values| {
+            let mut pairs = Vec::<(Self::Item, Self::Item)>::new();
+            pairs.reserve(nb);
+            for chunk in values.chunks(2) {
+                assert!(chunk.len() == 2);
+                // TODO: remove clones. http://stackoverflow.com/q/37097395/539465
+                pairs.push((chunk.get(0).unwrap().clone(), chunk.get(1).unwrap().clone()));
+            }
+            pairs
+        })
+    }
 }

Original file line number	Diff line number	Diff line change
`@@ -8,6 +8,7 @@ pub trait VarStack : Debug {`
`8`	`8`	`fn pop_many(&mut self, count: usize) -> Option<Vec<Self::Item>>;`
`9`	`9`	`fn push(&mut self, value: Self::Item);`
`10`	`10`	`fn pop_all_and_get_n_last(&mut self, nb: usize) -> Option<Vec<Self::Item>>;`
	`11`	`+ fn pop_n_pairs(&mut self, nb: usize) -> Option<Vec<(Self::Item, Self::Item)>>;`
`11`	`12`	`}`
`12`	`13`
`13`	`14`	`#[derive(Debug)]`
`@@ -27,7 +28,7 @@ impl<Item> VectorVarStack<Item> {`
`27`	`28`	`}`
`28`	`29`	`}`
`29`	`30`
`30`		`-impl<Item> VarStack for VectorVarStack<Item> where Item: Debug {`
	`31`	`+impl<Item: Clone> VarStack for VectorVarStack<Item> where Item: Debug {`
`31`	`32`	`type Item = Item;`
`32`	`33`
`33`	`34`	`fn top(&self) -> Option<&Self::Item> {`
`@@ -61,4 +62,17 @@ impl<Item> VarStack for VectorVarStack<Item> where Item: Debug {`
`61`	`62`	`Some(self.vector.drain(..).collect())`
`62`	`63`	`}`
`63`	`64`	`}`
	`65`	`+`
	`66`	`+ fn pop_n_pairs(&mut self, nb: usize) -> Option<Vec<(Self::Item, Self::Item)>> {`
	`67`	`+ self.pop_many(nb*2).map(\|values\| {`
	`68`	`+ let mut pairs = Vec::<(Self::Item, Self::Item)>::new();`
	`69`	`+ pairs.reserve(nb);`
	`70`	`+ for chunk in values.chunks(2) {`
	`71`	`+ assert!(chunk.len() == 2);`
	`72`	`+ // TODO: remove clones. http://stackoverflow.com/q/37097395/539465`
	`73`	`+ pairs.push((chunk.get(0).unwrap().clone(), chunk.get(1).unwrap().clone()));`
	`74`	`+ }`
	`75`	`+ pairs`
	`76`	`+ })`
	`77`	`+ }`
`64`	`78`	`}`