Support multiple context managers in a 'with' statement

This commit is contained in:
K. Lange 2021-03-18 21:25:19 +09:00
parent ee731257dd
commit 46dada92fc
4 changed files with 86 additions and 7 deletions

View File

@ -1435,6 +1435,7 @@ static void withStatement() {
/* We only need this for block() */ /* We only need this for block() */
size_t blockWidth = (parser.previous.type == TOKEN_INDENTATION) ? parser.previous.length : 0; size_t blockWidth = (parser.previous.type == TOKEN_INDENTATION) ? parser.previous.length : 0;
KrkToken myPrevious = parser.previous;
/* Collect the with token that started this statement */ /* Collect the with token that started this statement */
advance(); advance();
@ -1453,19 +1454,25 @@ static void withStatement() {
markInitialized(); markInitialized();
} }
consume(TOKEN_COLON, "Expected ':' after with statement");
/* Storage for return / exception */ /* Storage for return / exception */
addLocal(syntheticToken("")); addLocal(syntheticToken(""));
markInitialized();
/* Handler object */ /* Handler object */
addLocal(syntheticToken("")); addLocal(syntheticToken(""));
int withJump = emitJump(OP_PUSH_WITH); int withJump = emitJump(OP_PUSH_WITH);
markInitialized(); markInitialized();
beginScope(); if (check(TOKEN_COMMA)) {
block(blockWidth,"with"); parser.previous = myPrevious;
endScope(); withStatement(); /* Keep nesting */
} else {
consume(TOKEN_COLON, "Expected ',' or ':' after with statement");
beginScope();
block(blockWidth,"with");
endScope();
}
patchJump(withJump); patchJump(withJump);
emitByte(OP_CLEANUP_WITH); emitByte(OP_CLEANUP_WITH);

View File

@ -300,7 +300,10 @@ inline void krk_push(KrkValue value) {
inline KrkValue krk_pop() { inline KrkValue krk_pop() {
krk_currentThread.stackTop--; krk_currentThread.stackTop--;
if (unlikely(krk_currentThread.stackTop < krk_currentThread.stack)) { if (unlikely(krk_currentThread.stackTop < krk_currentThread.stack)) {
fprintf(stderr, "Fatal error: stack underflow detected in VM.\n"); fprintf(stderr, "Fatal error: stack underflow detected in VM (krk_currentThread.stackTop = %p, krk_currentThread.stack = %p)\n",
(void*)krk_currentThread.stackTop,
(void*)krk_currentThread.stack);
abort();
return NONE_VAL(); return NONE_VAL();
} }
return *krk_currentThread.stackTop; return *krk_currentThread.stackTop;
@ -738,8 +741,10 @@ _finishKwarg:
/* We can't have had any kwargs. */ /* We can't have had any kwargs. */
if ((size_t)argCount > potentialPositionalArgs && closure->function->collectsArguments) { if ((size_t)argCount > potentialPositionalArgs && closure->function->collectsArguments) {
krk_push(NONE_VAL()); krk_push(NONE_VAL()); krk_pop(); krk_pop(); krk_push(NONE_VAL()); krk_push(NONE_VAL()); krk_pop(); krk_pop();
startOfPositionals[offsetOfExtraArgs] = krk_list_of(argCount - potentialPositionalArgs, KrkValue tmp = krk_list_of(argCount - potentialPositionalArgs,
&startOfPositionals[potentialPositionalArgs], 0); &startOfPositionals[potentialPositionalArgs], 0);
startOfPositionals = &krk_currentThread.stackTop[-argCount];
startOfPositionals[offsetOfExtraArgs] = tmp;
argCount = closure->function->requiredArgs + 1; argCount = closure->function->requiredArgs + 1;
argCountX = argCount - 1; argCountX = argCount - 1;
while (krk_currentThread.stackTop > startOfPositionals + argCount) krk_pop(); while (krk_currentThread.stackTop > startOfPositionals + argCount) krk_pop();
@ -2004,6 +2009,7 @@ _resumeHook: (void)0;
krk_push(NONE_VAL()); krk_push(NONE_VAL());
krk_push(NONE_VAL()); krk_push(NONE_VAL());
krk_callSimple(OBJECT_VAL(type->_exit), 4, 0); krk_callSimple(OBJECT_VAL(type->_exit), 4, 0);
if (krk_currentThread.flags & KRK_THREAD_HAS_EXCEPTION) goto _finishException;
} }
if (AS_HANDLER(handler).type != OP_RETURN) break; if (AS_HANDLER(handler).type != OP_RETURN) break;
krk_pop(); /* handler */ krk_pop(); /* handler */

39
test/testMultipleWith.krk Normal file
View File

@ -0,0 +1,39 @@
class Context:
def __init__(self,name):
self.name = name
def __enter__(self):
print("Entering",self.name)
def __exit__(self, *args):
print("Exiting",self.name,"with",[type(x) for x in args])
def nested():
print("Before")
with Context('a') as a:
with Context('b') as b:
with Context('c') as c:
print("Hello, world.")
print("After")
nested()
def simple():
print("Before")
with Context('a') as a, Context('b') as b, Context('c') as c:
print("Hello, world.")
print("After")
simple()
def exceptional():
print("Before")
with Context('a') as a, Context('b') as b, Context('c') as c:
print("Prior")
raise ValueError("oh no")
print("Shouldn't happen")
print("After")
try:
exceptional()
except Exception as e:
print("Raised",repr(e))

View File

@ -0,0 +1,27 @@
Before
Entering a
Entering b
Entering c
Hello, world.
Exiting c with [<class 'NoneType'>, <class 'NoneType'>, <class 'NoneType'>]
Exiting b with [<class 'NoneType'>, <class 'NoneType'>, <class 'NoneType'>]
Exiting a with [<class 'NoneType'>, <class 'NoneType'>, <class 'NoneType'>]
After
Before
Entering a
Entering b
Entering c
Hello, world.
Exiting c with [<class 'NoneType'>, <class 'NoneType'>, <class 'NoneType'>]
Exiting b with [<class 'NoneType'>, <class 'NoneType'>, <class 'NoneType'>]
Exiting a with [<class 'NoneType'>, <class 'NoneType'>, <class 'NoneType'>]
After
Before
Entering a
Entering b
Entering c
Prior
Exiting c with [<class 'type'>, <class 'ValueError'>, <class 'list'>]
Exiting b with [<class 'type'>, <class 'ValueError'>, <class 'list'>]
Exiting a with [<class 'type'>, <class 'ValueError'>, <class 'list'>]
Raised ValueError('oh no')