Merge pull request #389 from pfalcon/with-statement

With statement implementation

7 files changed, 186 insertions, 2 deletions

diff --git a/py/compile.c b/py/compile.c
index 48c9c08826..0e33ca2688 100644
--- a/py/compile.c
+++ b/py/compile.c
@@ -1798,6 +1798,7 @@ void compile_with_stmt_helper(compiler_t *comp, int n, mp_parse_node_t *nodes, m
             EMIT_ARG(setup_with, l_end);
             EMIT(pop_top);
         }
+        compile_increase_except_level(comp);
         // compile additional pre-bits and the body
         compile_with_stmt_helper(comp, n - 1, nodes + 1, body);
         // finish this with block
@@ -1805,6 +1806,7 @@ void compile_with_stmt_helper(compiler_t *comp, int n, mp_parse_node_t *nodes, m
         EMIT_ARG(load_const_tok, MP_TOKEN_KW_NONE);
         EMIT_ARG(label_assign, l_end);
         EMIT(with_cleanup);
+        compile_decrease_except_level(comp);
         EMIT(end_finally);
     }
 }
diff --git a/py/qstrdefs.h b/py/qstrdefs.h
index 3be6295067..457043938b 100644
--- a/py/qstrdefs.h
+++ b/py/qstrdefs.h
@@ -16,6 +16,8 @@ Q(__qualname__)
 Q(__repl_print__)
 
 Q(__bool__)
+Q(__enter__)
+Q(__exit__)
 Q(__len__)
 Q(__iter__)
 Q(__getitem__)
diff --git a/py/vm.c b/py/vm.c
index 0ff27254a0..f097ff0a82 100644
--- a/py/vm.c
+++ b/py/vm.c
@@ -381,6 +381,73 @@ dispatch_loop:
                         break;
                         */
 
+                    case MP_BC_SETUP_WITH: {
+                        obj1 = TOP();
+                        SET_TOP(rt_load_attr(obj1, MP_QSTR___exit__));
+                        mp_obj_t dest[2];
+                        rt_load_method(obj1, MP_QSTR___enter__, dest);
+                        obj2 = rt_call_method_n_kw(0, 0, dest);
+                        SETUP_BLOCK();
+                        PUSH(obj2);
+                        break;
+                    }
+
+                    case MP_BC_WITH_CLEANUP: {
+                        static const mp_obj_t no_exc[] = {mp_const_none, mp_const_none, mp_const_none};
+                        if (TOP() == mp_const_none) {
+                            sp--;
+                            obj1 = TOP();
+                            SET_TOP(mp_const_none);
+                            obj2 = rt_call_function_n_kw(obj1, 3, 0, no_exc);
+                        } else if (MP_OBJ_IS_SMALL_INT(TOP())) {
+                            mp_obj_t cause = POP();
+                            switch (MP_OBJ_SMALL_INT_VALUE(cause)) {
+                                case UNWIND_RETURN: {
+                                    mp_obj_t retval = POP();
+                                    obj2 = rt_call_function_n_kw(TOP(), 3, 0, no_exc);
+                                    SET_TOP(retval);
+                                    PUSH(cause);
+                                    break;
+                                }
+                                case UNWIND_JUMP: {
+                                    obj2 = rt_call_function_n_kw(sp[-2], 3, 0, no_exc);
+                                    // Pop __exit__ boundmethod at sp[-2]
+                                    sp[-2] = sp[-1];
+                                    sp[-1] = sp[0];
+                                    SET_TOP(cause);
+                                    break;
+                                }
+                                default:
+                                    assert(0);
+                            }
+                        } else if (mp_obj_is_exception_type(TOP())) {
+                            mp_obj_t args[3] = {sp[0], sp[-1], sp[-2]};
+                            obj2 = rt_call_function_n_kw(sp[-3], 3, 0, args);
+                            // Pop __exit__ boundmethod at sp[-3]
+                            // TODO: Once semantics is proven, optimize for case when obj2 == True
+                            sp[-3] = sp[-2];
+                            sp[-2] = sp[-1];
+                            sp[-1] = sp[0];
+                            sp--;
+                            if (rt_is_true(obj2)) {
+                                // This is what CPython does
+                                //PUSH(MP_OBJ_NEW_SMALL_INT(UNWIND_SILENCED));
+                                // But what we need to do is - pop exception from value stack...
+                                sp -= 3;
+                                // ... pop with exception handler, and signal END_FINALLY
+                                // to just execute finally handler normally (signalled by None
+                                // on value stack)
+                                assert(exc_sp >= exc_stack);
+                                assert(exc_sp->opcode == MP_BC_SETUP_WITH);
+                                exc_sp--;
+                                PUSH(mp_const_none);
+                            }
+                        } else {
+                            assert(0);
+                        }
+                        break;
+                    }
+
                     case MP_BC_UNWIND_JUMP:
                         DECODE_SLABEL;
                         PUSH((void*)(ip + unum)); // push destination ip for jump
@@ -390,7 +457,7 @@ unwind_jump:
                         while (unum > 0) {
                             unum -= 1;
                             assert(exc_sp >= exc_stack);
-                            if (exc_sp->opcode == MP_BC_SETUP_FINALLY) {
+                            if (exc_sp->opcode == MP_BC_SETUP_FINALLY || exc_sp->opcode == MP_BC_SETUP_WITH) {
                                 // We're going to run "finally" code as a coroutine
                                 // (not calling it recursively). Set up a sentinel
                                 // on a stack so it can return back to us when it is
@@ -604,7 +671,7 @@ unwind_jump:
                     case MP_BC_RETURN_VALUE:
 unwind_return:
                         while (exc_sp >= exc_stack) {
-                            if (exc_sp->opcode == MP_BC_SETUP_FINALLY) {
+                            if (exc_sp->opcode == MP_BC_SETUP_FINALLY || exc_sp->opcode == MP_BC_SETUP_WITH) {
                                 // We're going to run "finally" code as a coroutine
                                 // (not calling it recursively). Set up a sentinel
                                 // on a stack so it can return back to us when it is
diff --git a/tests/basics/with-break.py b/tests/basics/with-break.py
new file mode 100644
index 0000000000..f1063d5826
--- /dev/null
+++ b/tests/basics/with-break.py
@@ -0,0 +1,14 @@
+class CtxMgr:
+
+    def __enter__(self):
+        print("__enter__")
+        return self
+
+    def __exit__(self, a, b, c):
+        print("__exit__", repr(a), repr(b))
+
+for i in range(5):
+    print(i)
+    with CtxMgr():
+        if i == 3:
+            break
diff --git a/tests/basics/with-continue.py b/tests/basics/with-continue.py
new file mode 100644
index 0000000000..fc2b24bd4b
--- /dev/null
+++ b/tests/basics/with-continue.py
@@ -0,0 +1,14 @@
+class CtxMgr:
+
+    def __enter__(self):
+        print("__enter__")
+        return self
+
+    def __exit__(self, a, b, c):
+        print("__exit__", repr(a), repr(b))
+
+for i in range(5):
+    print(i)
+    with CtxMgr():
+        if i == 3:
+            continue
diff --git a/tests/basics/with-return.py b/tests/basics/with-return.py
new file mode 100644
index 0000000000..cb0135c8b3
--- /dev/null
+++ b/tests/basics/with-return.py
@@ -0,0 +1,14 @@
+class CtxMgr:
+
+    def __enter__(self):
+        print("__enter__")
+        return self
+
+    def __exit__(self, a, b, c):
+        print("__exit__", repr(a), repr(b))
+
+def foo():
+    with CtxMgr():
+        return 4
+
+print(foo())
diff --git a/tests/basics/with1.py b/tests/basics/with1.py
new file mode 100644
index 0000000000..3db1d380d4
--- /dev/null
+++ b/tests/basics/with1.py
@@ -0,0 +1,71 @@
+class CtxMgr:
+
+    def __enter__(self):
+        print("__enter__")
+        return self
+
+    def __exit__(self, a, b, c):
+        print("__exit__", repr(a), repr(b))
+
+
+with CtxMgr() as a:
+    print(isinstance(a, CtxMgr))
+
+try:
+    with CtxMgr() as a:
+        raise ValueError
+except ValueError:
+    print("ValueError")
+
+
+class CtxMgr2:
+
+    def __enter__(self):
+        print("__enter__")
+        return self
+
+    def __exit__(self, a, b, c):
+        print("__exit__", repr(a), repr(b))
+        return True
+
+try:
+    with CtxMgr2() as a:
+        raise ValueError
+    print("No ValueError2")
+except ValueError:
+    print("ValueError2")
+
+
+# These recursive try-finally tests are attempt to get some interpretation
+# of last phrase in http://docs.python.org/3.4/library/dis.html#opcode-WITH_CLEANUP
+# "If the stack represents an exception, and the function call returns a ‘true’
+# value, this information is “zapped” and replaced with a single WHY_SILENCED
+# to prevent END_FINALLY from re-raising the exception. (But non-local gotos
+# will still be resumed.)"
+print("===")
+with CtxMgr2() as a:
+    try:
+        try:
+            raise ValueError
+            print("No ValueError3")
+        finally:
+            print("finally1")
+    finally:
+        print("finally2")
+
+print("===")
+try:
+    try:
+        with CtxMgr2() as a:
+            try:
+                try:
+                    raise ValueError
+                    print("No ValueError3")
+                finally:
+                    print("finally1")
+            finally:
+                print("finally2")
+    finally:
+        print("finally3")
+finally:
+    print("finally4")