Add tuple unpacking and 'if' conditions to list comprehensions

This commit is contained in:
K. Lange 2021-01-08 17:58:33 +09:00
parent fc05327c42
commit 070b4bc356
3 changed files with 40 additions and 5 deletions

View File

@ -1610,13 +1610,17 @@ static void list(int canAssign) {
/* x in... */
ssize_t loopInd = current->localCount;
varDeclaration();
defineVariable(loopInd);
ssize_t varCount = 0;
do {
varDeclaration();
defineVariable(loopInd);
varCount++;
} while (match(TOKEN_COMMA));
consume(TOKEN_IN, "Only iterator loops (for ... in ...) are allowed in list comprehensions.");
beginScope();
expression();
parsePrecedence(PREC_OR); /* Otherwise we can get trapped on a ternary */
endScope();
/* iterable... */
@ -1649,10 +1653,29 @@ static void list(int canAssign) {
* and there's no feasible way they can return themselves without
* our intended sentinel meaning, right? Surely? */
EMIT_CONSTANT_OP(OP_GET_LOCAL, indLoopIter);
emitBytes(OP_EQUAL, OP_NOT);
int exitJump = emitJump(OP_JUMP_IF_FALSE);
emitByte(OP_EQUAL);
int exitJump = emitJump(OP_JUMP_IF_TRUE);
emitByte(OP_POP);
/* Unpack tuple */
if (varCount > 1) {
EMIT_CONSTANT_OP(OP_GET_LOCAL, loopInd);
EMIT_CONSTANT_OP(OP_UNPACK_TUPLE, varCount);
for (ssize_t i = loopInd + varCount - 1; i >= loopInd; i--) {
EMIT_CONSTANT_OP(OP_SET_LOCAL, i);
emitByte(OP_POP);
}
}
if (match(TOKEN_IF)) {
parsePrecedence(PREC_OR);
int acceptJump = emitJump(OP_JUMP_IF_TRUE);
emitByte(OP_POP); /* Pop condition */
emitLoop(loopStart);
patchJump(acceptJump);
emitByte(OP_POP); /* Pop condition */
}
/* Now we can rewind the scanner to have it parse the original
* expression that uses our iterated values! */
KrkScanner scannerAfter = krk_tellScanner();

View File

@ -29,3 +29,12 @@ for i in {'a': 1, 2: 'b'}.items():
print(i)
# ('a', 1)
# (2, 'b')
print(['{}: {}'.format(repr(k),repr(v)) for k, v in {1: 2, 3: 4}.items()])
# ['1: 2', '3: 4']
print([x for x,y in [(1,2),(3,4),(5,6)] if x + y > 3])
# [3, 5]
print(', '.join(['{}: {}'.format(repr(k),repr(v)) for k, v in {1: 2, '3': 4, 5: 'six', 'seven': 8}.items() if isinstance(k,str)]))
# '3': 4, 'seven': 8

View File

@ -9,3 +9,6 @@ a 1
2 b
('a', 1)
(2, 'b')
['1: 2', '3: 4']
[3, 5]
'3': 4, 'seven': 8