Merge pull request #49 from kmike/coroutine-bug

fixed an error with callbacks returned from wrapped Lua coroutines

lupa/_lupa.pyx

@@ -872,8 +872,10 @@ cdef _LuaObject new_lua_thread_or_function(LuaRuntime runtime, lua_State* L, int
 
 cdef object resume_lua_thread(_LuaThread thread, tuple args):
     cdef lua_State* co = thread._co_state
-    cdef int result, i, nargs = 0
+    cdef lua_State* L = thread._state
+    cdef int status, i, nargs = 0, nres = 0
     lock_runtime(thread._runtime)
+    old_top = lua.lua_gettop(L)
     try:
         if lua.lua_status(co) == 0 and lua.lua_gettop(co) == 0:
             # already terminated
@@ -882,18 +884,25 @@ cdef object resume_lua_thread(_LuaThread thread, tuple args):
             nargs = len(args)
             push_lua_arguments(thread._runtime, co, args)
         with nogil:
-            result = lua.lua_resume(co, nargs)
-        if result != lua.LUA_YIELD:
-            if result == 0:
+            status = lua.lua_resume(co, L, nargs)
+        nres = lua.lua_gettop(co)
+        if status != lua.LUA_YIELD:
+            if status == 0:
                 # terminated
-                if lua.lua_gettop(co) == 0:
+                if nres == 0:
                     # no values left to return
                     raise StopIteration
             else:
-                raise_lua_error(thread._runtime, co, result)
-        return unpack_lua_results(thread._runtime, co)
+                raise_lua_error(thread._runtime, co, status)
+
+        # Move yielded values to the main state before unpacking.
+        # This is what Lua's internal auxresume function is doing;
+        # it affects wrapped Lua functions returned to Python.
+        lua.lua_xmove(co, L, nres)
+        return unpack_lua_results(thread._runtime, L)
     finally:
-        lua.lua_settop(co, 0)  # FIXME?
+        # FIXME: check that coroutine state is OK in case of errors?
+        lua.lua_settop(L, old_top)
         unlock_runtime(thread._runtime)
 
 

lupa/lua.pxd

@@ -150,7 +150,7 @@ cdef extern from "lua.h" nogil:
 
     # coroutine functions
     int  lua_yield (lua_State *L, int nresults)
-    int  lua_resume "__lupa_lua_resume" (lua_State *L, int narg)
+    int  lua_resume "__lupa_lua_resume" (lua_State *L, lua_State *from_, int narg)
     int  lua_status (lua_State *L)
 
     # garbage-collection function and options

lupa/lupa_defs.h

@@ -3,12 +3,12 @@
  */
 
 #if LUA_VERSION_NUM >= 502
-#define __lupa_lua_resume(L, nargs)   lua_resume(L, NULL, nargs)
-#define lua_objlen(L, i)              lua_rawlen(L, (i))
+#define __lupa_lua_resume(L, from_, nargs)   lua_resume(L, from_, nargs)
+#define lua_objlen(L, i)                     lua_rawlen(L, (i))
 
 #else
 #if LUA_VERSION_NUM >= 501
-#define __lupa_lua_resume(L, nargs)   lua_resume(L, nargs)
+#define __lupa_lua_resume(L, from_, nargs)   lua_resume(L, nargs)
 
 #else
 #error Lupa requires at least Lua 5.1 or LuaJIT 2.x

lupa/tests/test.py

@@ -1372,6 +1372,101 @@ def test_coroutine_while_status(self):
         self.assertEqual([0,1,0,1,0,1], result)
 
 
+class TestLuaCoroutinesWithDebugHooks(SetupLuaRuntimeMixin, unittest.TestCase):
+
+    def _enable_hook(self):
+        self.lua.execute('''
+            steps = 0
+            debug.sethook(function () steps = steps + 1 end, '', 1)
+        ''')
+
+    def test_coroutine_yields_callback_debug_hook(self):
+        self.lua.execute('''
+            func = function()
+                coroutine.yield(function() return 123 end)
+            end
+        ''')
+        def _check():
+            coro = self.lua.eval('func').coroutine()
+            cb = next(coro)
+            self.assertEqual(cb(), 123)
+
+        # yielding a callback should work without a debug hook
+        _check()
+
+        # it should keep working after a debug hook is added
+        self._enable_hook()
+        _check()
+
+    def test_coroutine_yields_callback_debug_hook_nowrap(self):
+        resume = self.lua.eval("coroutine.resume")
+        self.lua.execute('''
+            func = function()
+                coroutine.yield(function() return 123 end)
+            end
+        ''')
+        def _check():
+            coro = self.lua.eval('func').coroutine()
+            ok, cb = resume(coro)
+            self.assertEqual(ok, True)
+            self.assertEqual(cb(), 123)
+
+        # yielding a callback should work without a debug hook
+        _check()
+
+        # it should keep working after a debug hook is added
+        self._enable_hook()
+        _check()
+
+    def test_coroutine_sets_callback_debug_hook(self):
+        self.lua.execute('''
+            func = function(dct)
+                dct['cb'] = function() return 123 end
+                coroutine.yield()
+            end
+        ''')
+        def _check(dct):
+            coro = self.lua.eval('func').coroutine(dct)
+            next(coro)
+            cb = dct['cb']
+            self.assertEqual(cb(), 123)
+
+        # sending a callback should work without a debug hook
+        _check({})
+
+        # enable debug hook and try again
+        self._enable_hook()
+
+        # it works with a Lua table wrapper
+        _check(self.lua.table_from({}))
+
+        # FIXME: but it fails with a regular dict
+        # _check({})
+
+    def test_coroutine_sets_callback_debug_hook_nowrap(self):
+        resume = self.lua.eval("coroutine.resume")
+        self.lua.execute('''
+            func = function(dct)
+                dct['cb'] = function() return 123 end
+                coroutine.yield()
+            end
+        ''')
+        def _check():
+            dct = {}
+            coro = self.lua.eval('func').coroutine()
+            resume(coro, dct)  # send initial value
+            resume(coro)
+            cb = dct['cb']
+            self.assertEqual(cb(), 123)
+
+        # sending a callback should work without a debug hook
+        _check()
+
+        # enable debug hook and try again
+        self._enable_hook()
+        _check()
+
+
 class TestLuaApplications(unittest.TestCase):
     def tearDown(self):
         gc.collect()