Skip to content

Commit 1f30693

Browse files
windelbouwmanwindelbouwman
committed
Allow unpack in assignment of list as well as trailing comma for 1 sized tuples
1 parent 57dbfd1 commit 1f30693

File tree

7 files changed

+119
-81
lines changed

7 files changed

+119
-81
lines changed

‎parser/src/ast.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -216,8 +216,8 @@ pub enum Expression {
216216
pub struct Parameters {
217217
pub args: Vec<String>,
218218
pub kwonlyargs: Vec<String>,
219-
pub vararg: Option<String>,
220-
pub kwarg: Option<String>,
219+
pub vararg: Option<Option<String>>, // Optionally we handle optionally named '*args' or '*'
220+
pub kwarg: Option<Option<String>>,
221221
pub defaults: Vec<Expression>,
222222
pub kw_defaults: Vec<Option<Expression>>,
223223
}

‎parser/src/python.lalrpop‎

Lines changed: 57 additions & 67 deletions
Original file line numberDiff line numberDiff line change
@@ -76,13 +76,6 @@ DelStatement: ast::LocatedStatement = {
7676

7777
ExpressionStatement: ast::LocatedStatement = {
7878
<loc:@L> <expr:TestOrStarExprList> <suffix:AssignSuffix*> => {
79-
// First build tuple from first item:
80-
let expr = if expr.len() > 1 {
81-
ast::Expression::Tuple { elements: expr }
82-
} else {
83-
expr.into_iter().next().unwrap()
84-
};
85-
8679
// Just an expression, no assignment:
8780
if suffix.is_empty() {
8881
ast::LocatedStatement {
@@ -106,12 +99,6 @@ ExpressionStatement: ast::LocatedStatement = {
10699
}
107100
},
108101
<loc:@L> <expr:TestOrStarExprList> <op:AugAssign> <e2:TestList> => {
109-
let expr = if expr.len() > 1 {
110-
ast::Expression::Tuple { elements: expr }
111-
} else {
112-
expr.into_iter().next().unwrap()
113-
};
114-
115102
// TODO: this works in most cases:
116103
let rhs = e2.into_iter().next().unwrap();
117104
ast::LocatedStatement {
@@ -134,11 +121,19 @@ AssignSuffix: ast::Expression = {
134121
"=" <e:YieldExpr> => e,
135122
};
136123

137-
TestOrStarExprList: Vec<ast::Expression> = {
138-
<e:TestOrStarExpr> <e2:("," TestOrStarExpr)*> => {
124+
TestOrStarExprList: ast::Expression = {
125+
<e:TestOrStarExpr> <e2:("," TestOrStarExpr)*> <comma:","?> => {
139126
let mut res = vec![e];
140127
res.extend(e2.into_iter().map(|x| x.1));
141-
res
128+
129+
// First build tuple from first item:
130+
let expr = if (res.len() > 1) || comma.is_some() {
131+
ast::Expression::Tuple { elements: res }
132+
} else {
133+
res.into_iter().next().unwrap()
134+
};
135+
136+
expr
142137
}
143138
};
144139

@@ -453,29 +448,8 @@ Parameters: ast::Parameters = {
453448
// parameters are (String, None), kwargs are (String, Some(Test)) where Test is
454449
// the default
455450
TypedArgsList: ast::Parameters = {
456-
<param1:TypedParameterDef> <param2:("," TypedParameterDef)*> <args2:("," ParameterListStarArgs)?> => {
457-
// Combine first parameters:
458-
let mut args = vec![param1];
459-
args.extend(param2.into_iter().map(|x| x.1));
460-
461-
let mut names = vec![];
462-
let mut default_elements = vec![];
463-
464-
for (name, default) in args.into_iter() {
465-
names.push(name.clone());
466-
if let Some(default) = default {
467-
default_elements.push(default);
468-
} else {
469-
if default_elements.len() > 0 {
470-
// Once we have started with defaults, all remaining arguments must
471-
// have defaults
472-
panic!(
473-
"non-default argument follows default argument: {}",
474-
name
475-
);
476-
}
477-
}
478-
}
451+
<param1:TypedParameters> <args2:("," ParameterListStarArgs)?> => {
452+
let (names, default_elements) = param1;
479453

480454
// Now gather rest of parameters:
481455
let (vararg, kwonlyargs, kw_defaults, kwarg) = match args2 {
@@ -492,35 +466,14 @@ TypedArgsList: ast::Parameters = {
492466
kw_defaults: kw_defaults,
493467
}
494468
},
495-
<param1:TypedParameterDef> <param2:("," TypedParameterDef)*> <kw:("," "**" Identifier)> => {
496-
// Combine first parameters:
497-
let mut args = vec![param1];
498-
args.extend(param2.into_iter().map(|x| x.1));
499-
500-
let mut names = vec![];
501-
let mut default_elements = vec![];
502-
503-
for (name, default) in args.into_iter() {
504-
names.push(name.clone());
505-
if let Some(default) = default {
506-
default_elements.push(default);
507-
} else {
508-
if default_elements.len() > 0 {
509-
// Once we have started with defaults, all remaining arguments must
510-
// have defaults
511-
panic!(
512-
"non-default argument follows default argument: {}",
513-
name
514-
);
515-
}
516-
}
517-
}
469+
<param1:TypedParameters> <kw:("," KwargParameter)> => {
470+
let (names, default_elements) = param1;
518471

519472
// Now gather rest of parameters:
520473
let vararg = None;
521474
let kwonlyargs = vec![];
522475
let kw_defaults = vec![];
523-
let kwarg = Some(kw.2);
476+
let kwarg = Some(kw.1);
524477

525478
ast::Parameters {
526479
args: names,
@@ -542,7 +495,7 @@ TypedArgsList: ast::Parameters = {
542495
kw_defaults: kw_defaults,
543496
}
544497
},
545-
"**" <kw:Identifier> => {
498+
<kw:KwargParameter> => {
546499
ast::Parameters {
547500
args: vec![],
548501
kwonlyargs: vec![],
@@ -554,6 +507,37 @@ TypedArgsList: ast::Parameters = {
554507
},
555508
};
556509

510+
// Use inline here to make sure the "," is not creating an ambiguity.
511+
#[inline]
512+
TypedParameters: (Vec<String>, Vec<ast::Expression>) = {
513+
<param1:TypedParameterDef> <param2:("," TypedParameterDef)*> => {
514+
// Combine first parameters:
515+
let mut args = vec![param1];
516+
args.extend(param2.into_iter().map(|x| x.1));
517+
518+
let mut names = vec![];
519+
let mut default_elements = vec![];
520+
521+
for (name, default) in args.into_iter() {
522+
names.push(name.clone());
523+
if let Some(default) = default {
524+
default_elements.push(default);
525+
} else {
526+
if default_elements.len() > 0 {
527+
// Once we have started with defaults, all remaining arguments must
528+
// have defaults
529+
panic!(
530+
"non-default argument follows default argument: {}",
531+
name
532+
);
533+
}
534+
}
535+
}
536+
537+
(names, default_elements)
538+
}
539+
};
540+
557541
TypedParameterDef: (String, Option<ast::Expression>) = {
558542
<i:TypedParameter> => (i, None),
559543
<i:TypedParameter> "=" <e:Test> => (i, Some(e)),
@@ -564,8 +548,8 @@ TypedParameter: String = {
564548
Identifier,
565549
};
566550

567-
ParameterListStarArgs: (Option<String>, Vec<String>, Vec<Option<ast::Expression>>, Option<String>) = {
568-
"*" <va:Identifier> <kw:("," TypedParameterDef)*> <kwarg:("," "**" Identifier)?> => {
551+
ParameterListStarArgs: (Option<Option<String>>, Vec<String>, Vec<Option<ast::Expression>>, Option<Option<String>>) = {
552+
"*" <va:Identifier?> <kw:("," TypedParameterDef)*> <kwarg:("," KwargParameter)?> => {
569553
// Extract keyword arguments:
570554
let mut kwonlyargs = vec![];
571555
let mut kw_defaults = vec![];
@@ -575,14 +559,20 @@ ParameterListStarArgs: (Option<String>, Vec<String>, Vec<Option<ast::Expression>
575559
}
576560

577561
let kwarg = match kwarg {
578-
Some((_, _, name)) => Some(name),
562+
Some((_, name)) => Some(name),
579563
None => None,
580564
};
581565

582566
(Some(va), kwonlyargs, kw_defaults, kwarg)
583567
}
584568
};
585569

570+
KwargParameter: Option<String> = {
571+
"**" <kwarg:Identifier?> => {
572+
kwarg
573+
}
574+
};
575+
586576
ClassDef: ast::LocatedStatement = {
587577
<d:Decorator*> <loc:@L> "class" <n:Identifier> <a:("(" ArgumentList ")")?> ":" <s:Suite> => {
588578
let (bases, keywords) = match a {

‎tests/snippets/assignment.py‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,3 +40,14 @@
4040
assert b == []
4141
assert c == 2
4242
assert d == 3
43+
44+
a, = [1]
45+
assert a == 1
46+
47+
def g():
48+
yield 1337
49+
yield 42
50+
51+
a, b = g()
52+
assert a == 1337
53+
assert b == 42

‎tests/snippets/function_args.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,3 +32,9 @@ def va2(*args, **kwargs):
3232
x = (5, 4)
3333
# TODO:
3434
# va2(*x)
35+
36+
# def va3(x, *, b=2):
37+
# pass
38+
39+
# va3(1, 2, 3, b=10)
40+

‎vm/src/bytecode.rs‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -21,10 +21,10 @@ pub struct CodeObject {
2121
pub instructions: Vec<Instruction>,
2222
pub label_map: HashMap<Label, usize>,
2323
pub locations: Vec<ast::Location>,
24-
pub arg_names: Vec<String>, // Names of positional arguments
25-
pub varargs: Option<String>, // *args
24+
pub arg_names: Vec<String>, // Names of positional arguments
25+
pub varargs: Option<Option<String>>, // *args or *
2626
pub kwonlyarg_names: Vec<String>,
27-
pub varkeywords: Option<String>, // **kwargs
27+
pub varkeywords: Option<Option<String>>, // **kwargs or **
2828
pub source_path: Option<String>,
2929
pub obj_name: String, // Name of the object that created this code object
3030
pub is_generator: bool,
@@ -33,9 +33,9 @@ pub struct CodeObject {
3333
impl CodeObject {
3434
pub fn new(
3535
arg_names: Vec<String>,
36-
varargs: Option<String>,
36+
varargs: Option<Option<String>>,
3737
kwonlyarg_names: Vec<String>,
38-
varkeywords: Option<String>,
38+
varkeywords: Option<Option<String>>,
3939
source_path: Option<String>,
4040
obj_name: String,
4141
) -> CodeObject {

‎vm/src/frame.rs‎

Lines changed: 25 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -625,8 +625,11 @@ impl Frame {
625625
}
626626
bytecode::Instruction::UnpackSequence { size } => {
627627
let value = self.pop_value();
628+
let elements = match self.extract_elements(vm, &value) {
629+
Ok(elements) => elements,
630+
Err(err) => return Some(Err(err)),
631+
};
628632

629-
let elements = objtuple::get_elements(&value);
630633
if elements.len() != *size {
631634
Some(Err(vm.new_value_error(
632635
"Wrong number of values to unpack".to_string(),
@@ -640,8 +643,10 @@ impl Frame {
640643
}
641644
bytecode::Instruction::UnpackEx { before, after } => {
642645
let value = self.pop_value();
643-
644-
let elements = objtuple::get_elements(&value);
646+
let elements = match self.extract_elements(vm, &value) {
647+
Ok(elements) => elements,
648+
Err(err) => return Some(Err(err)),
649+
};
645650
let min_expected = *before + *after;
646651
if elements.len() < min_expected {
647652
Some(Err(vm.new_value_error(format!(
@@ -687,6 +692,23 @@ impl Frame {
687692
}
688693
}
689694

695+
fn extract_elements(
696+
&mut self,
697+
vm: &mut VirtualMachine,
698+
value: &PyObjectRef,
699+
) -> Result<Vec<PyObjectRef>, PyObjectRef> {
700+
// Extract elements from item, if possible:
701+
let elements = if objtype::isinstance(value, &vm.ctx.tuple_type()) {
702+
objtuple::get_elements(value)
703+
} else if objtype::isinstance(value, &vm.ctx.list_type()) {
704+
objlist::get_elements(value)
705+
} else {
706+
let iter = objiter::get_iter(vm, value)?;
707+
objiter::get_all(vm, &iter)?
708+
};
709+
Ok(elements)
710+
}
711+
690712
fn import(
691713
&mut self,
692714
vm: &mut VirtualMachine,

‎vm/src/vm.rs‎

Lines changed: 13 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -235,14 +235,18 @@ impl VirtualMachine {
235235
}
236236

237237
// Pack other positional arguments in to *args:
238-
if let Some(vararg_name) = &code_object.varargs {
238+
if let Some(vararg) = &code_object.varargs {
239239
let mut last_args = vec![];
240240
for i in n..nargs {
241241
let arg = &args.args[i];
242242
last_args.push(arg.clone());
243243
}
244244
let vararg_value = self.ctx.new_tuple(last_args);
245-
scope.set_item(vararg_name, vararg_value);
245+
246+
// If we have a name (not '*' only) then store it:
247+
if let Some(vararg_name) = vararg {
248+
scope.set_item(vararg_name, vararg_value);
249+
}
246250
} else {
247251
// Check the number of positional arguments
248252
if nargs > nexpected_args {
@@ -254,9 +258,14 @@ impl VirtualMachine {
254258
}
255259

256260
// Do we support `**kwargs` ?
257-
let kwargs = if let Some(name) = &code_object.varkeywords {
261+
let kwargs = if let Some(kwargs) = &code_object.varkeywords {
258262
let d = self.new_dict();
259-
scope.set_item(&name, d.clone());
263+
264+
// Store when we have a name:
265+
if let Some(kwargs_name) = kwargs {
266+
scope.set_item(&kwargs_name, d.clone());
267+
}
268+
260269
Some(d)
261270
} else {
262271
None

0 commit comments

Comments
 (0)