Skip to content

Commit 3f6c716

Browse files
windelbouwmanwindelbouwman
committed
Implement star expand argument in function call.
1 parent 851d263 commit 3f6c716

File tree

4 files changed

+113
-69
lines changed

4 files changed

+113
-69
lines changed

‎tests/snippets/function_args.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -27,12 +27,13 @@ def va(a, b=2, *c, d, **e):
2727

2828
def va2(*args, **kwargs):
2929
assert args == (5, 4)
30+
assert len(kwargs) == 0
3031

3132
va2(5, 4)
3233
x = (5, 4)
33-
# TODO:
34-
# va2(*x)
34+
va2(*x)
3535

36+
va2(5, *x[1:])
3637
# def va3(x, *, b=2):
3738
# pass
3839

‎vm/src/bytecode.rs

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -122,10 +122,7 @@ pub enum Instruction {
122122
flags: FunctionOpArg,
123123
},
124124
CallFunction {
125-
count: usize,
126-
},
127-
CallFunctionKw {
128-
count: usize,
125+
typ: CallType,
129126
},
130127
ForIter,
131128
ReturnValue,
@@ -158,6 +155,7 @@ pub enum Instruction {
158155
},
159156
BuildSet {
160157
size: usize,
158+
unpack: bool,
161159
},
162160
BuildMap {
163161
size: usize,
@@ -187,6 +185,13 @@ pub enum Instruction {
187185
Unpack,
188186
}
189187

188+
#[derive(Debug, Clone, PartialEq)]
189+
pub enum CallType {
190+
Positional(usize),
191+
Keyword(usize),
192+
Ex(bool),
193+
}
194+
190195
#[derive(Debug, Clone, PartialEq)]
191196
pub enum Constant {
192197
Integer { value: i32 }, // TODO: replace by arbitrary big int math.

‎vm/src/compile.rs

Lines changed: 67 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99
extern crate rustpython_parser;
1010

1111
use self::rustpython_parser::{ast, parser};
12-
use super::bytecode::{self, CodeObject, Instruction};
12+
use super::bytecode::{self, CallType, CodeObject, Instruction};
1313
use super::pyobject::{PyObject, PyObjectKind, PyResult};
1414
use super::vm::VirtualMachine;
1515

@@ -330,7 +330,9 @@ impl Compiler {
330330
});
331331
self.emit(Instruction::Rotate { amount: 2 });
332332
self.compile_expression(exc_type)?;
333-
self.emit(Instruction::CallFunction { count: 2 });
333+
self.emit(Instruction::CallFunction {
334+
typ: CallType::Positional(2),
335+
});
334336

335337
// We cannot handle this exception type:
336338
self.emit(Instruction::JumpIfFalse {
@@ -478,7 +480,7 @@ impl Compiler {
478480
self.compile_expression(base)?;
479481
}
480482
self.emit(Instruction::CallFunction {
481-
count: 2 + bases.len(),
483+
typ: CallType::Positional(2 + bases.len()),
482484
});
483485

484486
self.apply_decorators(decorator_list);
@@ -498,10 +500,14 @@ impl Compiler {
498500
match msg {
499501
Some(e) => {
500502
self.compile_expression(e)?;
501-
self.emit(Instruction::CallFunction { count: 1 });
503+
self.emit(Instruction::CallFunction {
504+
typ: CallType::Positional(1),
505+
});
502506
}
503507
None => {
504-
self.emit(Instruction::CallFunction { count: 0 });
508+
self.emit(Instruction::CallFunction {
509+
typ: CallType::Positional(0),
510+
});
505511
}
506512
}
507513
self.emit(Instruction::Raise { argc: 1 });
@@ -633,7 +639,9 @@ impl Compiler {
633639
fn apply_decorators(&mut self, decorator_list: &Vec<ast::Expression>) {
634640
// Apply decorators:
635641
for _ in decorator_list {
636-
self.emit(Instruction::CallFunction { count: 1 });
642+
self.emit(Instruction::CallFunction {
643+
typ: CallType::Positional(1),
644+
});
637645
}
638646
}
639647

@@ -859,10 +867,11 @@ impl Compiler {
859867
}
860868
ast::Expression::Set { elements } => {
861869
let size = elements.len();
862-
for element in elements {
863-
self.compile_expression(element)?;
864-
}
865-
self.emit(Instruction::BuildSet { size: size });
870+
let must_unpack = self.gather_elements(elements)?;
871+
self.emit(Instruction::BuildSet {
872+
size: size,
873+
unpack: must_unpack,
874+
});
866875
}
867876
ast::Expression::Dict { elements } => {
868877
let size = elements.len();
@@ -968,37 +977,53 @@ impl Compiler {
968977
args: &Vec<ast::Expression>,
969978
keywords: &Vec<ast::Keyword>,
970979
) -> Result<(), String> {
971-
self.compile_expression(&*function)?;
980+
self.compile_expression(function)?;
972981
let count = args.len() + keywords.len();
973982

974983
// Normal arguments:
975-
for value in args {
976-
self.compile_expression(value)?;
977-
}
978-
979-
// Keyword arguments:
980-
if keywords.len() > 0 {
981-
let mut kwarg_names = vec![];
982-
for keyword in keywords {
983-
if let Some(name) = &keyword.name {
984-
kwarg_names.push(bytecode::Constant::String {
985-
value: name.to_string(),
986-
});
987-
} else {
988-
// This means **kwargs!
989-
panic!("name must be set");
990-
}
991-
self.compile_expression(&keyword.value)?;
992-
}
984+
let must_unpack = self.gather_elements(args)?;
993985

994-
self.emit(Instruction::LoadConst {
995-
value: bytecode::Constant::Tuple {
996-
elements: kwarg_names,
997-
},
986+
if must_unpack {
987+
self.emit(Instruction::BuildTuple {
988+
size: args.len(),
989+
unpack: true,
998990
});
999-
self.emit(Instruction::CallFunctionKw { count });
991+
if keywords.len() > 0 {
992+
unimplemented!()
993+
} else {
994+
self.emit(Instruction::CallFunction {
995+
typ: CallType::Ex(false),
996+
});
997+
}
1000998
} else {
1001-
self.emit(Instruction::CallFunction { count });
999+
// Keyword arguments:
1000+
if keywords.len() > 0 {
1001+
let mut kwarg_names = vec![];
1002+
for keyword in keywords {
1003+
if let Some(name) = &keyword.name {
1004+
kwarg_names.push(bytecode::Constant::String {
1005+
value: name.to_string(),
1006+
});
1007+
} else {
1008+
// This means **kwargs!
1009+
panic!("name must be set");
1010+
}
1011+
self.compile_expression(&keyword.value)?;
1012+
}
1013+
1014+
self.emit(Instruction::LoadConst {
1015+
value: bytecode::Constant::Tuple {
1016+
elements: kwarg_names,
1017+
},
1018+
});
1019+
self.emit(Instruction::CallFunction {
1020+
typ: CallType::Keyword(count),
1021+
});
1022+
} else {
1023+
self.emit(Instruction::CallFunction {
1024+
typ: CallType::Positional(count),
1025+
});
1026+
}
10021027
}
10031028
Ok(())
10041029
}
@@ -1068,7 +1093,10 @@ impl Compiler {
10681093
});
10691094
}
10701095
ast::ComprehensionKind::Set { .. } => {
1071-
self.emit(Instruction::BuildSet { size: 0 });
1096+
self.emit(Instruction::BuildSet {
1097+
size: 0,
1098+
unpack: false,
1099+
});
10721100
}
10731101
ast::ComprehensionKind::Dict { .. } => {
10741102
self.emit(Instruction::BuildMap { size: 0 });
@@ -1183,7 +1211,9 @@ impl Compiler {
11831211
self.emit(Instruction::GetIter);
11841212

11851213
// Call just created <listcomp> function:
1186-
self.emit(Instruction::CallFunction { count: 1 });
1214+
self.emit(Instruction::CallFunction {
1215+
typ: CallType::Positional(1),
1216+
});
11871217
Ok(())
11881218
}
11891219

‎vm/src/frame.rs

Lines changed: 34 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -233,8 +233,8 @@ impl Frame {
233233
self.push_value(list_obj);
234234
Ok(None)
235235
}
236-
bytecode::Instruction::BuildSet { size } => {
237-
let elements = self.pop_multiple(*size);
236+
bytecode::Instruction::BuildSet { size, unpack } => {
237+
let elements = self.get_elements(vm, *size, *unpack)?;
238238
let py_obj = vm.ctx.new_set(elements);
239239
self.push_value(py_obj);
240240
Ok(None)
@@ -444,34 +444,42 @@ impl Frame {
444444
self.push_value(obj);
445445
Ok(None)
446446
}
447-
bytecode::Instruction::CallFunction { count } => {
448-
let args: Vec<PyObjectRef> = self.pop_multiple(*count);
449-
let args = PyFuncArgs {
450-
args: args,
451-
kwargs: vec![],
447+
bytecode::Instruction::CallFunction { typ } => {
448+
let args = match typ {
449+
bytecode::CallType::Positional(count) => {
450+
let args: Vec<PyObjectRef> = self.pop_multiple(*count);
451+
PyFuncArgs {
452+
args: args,
453+
kwargs: vec![],
454+
}
455+
}
456+
bytecode::CallType::Keyword(count) => {
457+
let kwarg_names = self.pop_value();
458+
let args: Vec<PyObjectRef> = self.pop_multiple(*count);
459+
460+
let kwarg_names = kwarg_names
461+
.to_vec()
462+
.unwrap()
463+
.iter()
464+
.map(|pyobj| objstr::get_value(pyobj))
465+
.collect();
466+
PyFuncArgs::new(args, kwarg_names)
467+
}
468+
bytecode::CallType::Ex(has_kwargs) => {
469+
let kwarg_names = if *has_kwargs {
470+
self.pop_value();
471+
unimplemented!();
472+
} else {
473+
vec![]
474+
};
475+
let args = self.pop_value();
476+
let args = self.extract_elements(vm, &args)?;
477+
PyFuncArgs::new(args, kwarg_names)
478+
}
452479
};
453-
let func_ref = self.pop_value();
454480

455481
// Call function:
456-
let value = vm.invoke(func_ref, args)?;
457-
458-
self.push_value(value);
459-
Ok(None)
460-
}
461-
bytecode::Instruction::CallFunctionKw { count } => {
462-
let kwarg_names = self.pop_value();
463-
let args: Vec<PyObjectRef> = self.pop_multiple(*count);
464-
465-
let kwarg_names = kwarg_names
466-
.to_vec()
467-
.unwrap()
468-
.iter()
469-
.map(|pyobj| objstr::get_value(pyobj))
470-
.collect();
471-
let args = PyFuncArgs::new(args, kwarg_names);
472482
let func_ref = self.pop_value();
473-
474-
// Call function:
475483
let value = vm.invoke(func_ref, args)?;
476484
self.push_value(value);
477485
Ok(None)

0 commit comments

Comments
 (0)