Report stacktraces for remote_exec() (#101)

Author: João Eiras

execnet/gateway.py

@@ -7,7 +7,6 @@
 import os
 import inspect
 import types
-import linecache
 import textwrap
 import execnet
 from execnet.gateway_base import Message
@@ -111,22 +110,27 @@ def remote_exec(self, source, **kwargs):
             executing code.
         """
         call_name = None
+        file_name = None
         if isinstance(source, types.ModuleType):
-            linecache.updatecache(inspect.getsourcefile(source))
-            source = inspect.getsource(source)
+            file_name = inspect.getsourcefile(source)
+            if not file_name:
+                source = inspect.getsource(source)
+            else:
+                source = None
         elif isinstance(source, types.FunctionType):
             call_name = source.__name__
+            file_name = inspect.getsourcefile(source)
             source = _source_of_function(source)
         else:
             source = textwrap.dedent(str(source))
 
-        if call_name is None and kwargs:
+        if not call_name and kwargs:
             raise TypeError("can't pass kwargs to non-function remote_exec")
 
         channel = self.newchannel()
         self._send(Message.CHANNEL_EXEC,
                    channel.id,
-                   gateway_base.dumps_internal((source, call_name, kwargs)))
+                   gateway_base.dumps_internal((source, file_name, call_name, kwargs)))
         return channel
 
     def remote_init_threads(self, num=None):
@@ -186,7 +190,7 @@ def _source_of_function(function):
         args = inspect.getargspec(function)[0]
     else:
         args = sig.args
-    if args[0] != 'channel':
+    if not args or args[0] != 'channel':
         raise ValueError('expected first function argument to be `channel`')
 
     if gateway_base.ISPY3:
@@ -213,4 +217,5 @@ def _source_of_function(function):
             used_globals,
         )
 
-    return source
+    leading_ws = "\n" * (codeobj.co_firstlineno - 1)
+    return leading_ws + source

execnet/gateway_base.py

@@ -13,6 +13,7 @@
 """
 from __future__ import with_statement
 import sys
+import linecache
 import os
 import weakref
 import traceback
@@ -25,8 +26,9 @@
 ISPY3 = sys.version_info >= (3, 0)
 if ISPY3:
     from io import BytesIO
-    exec("def do_exec(co, loc): exec(co, loc)\n"
-         "def reraise(cls, val, tb): raise val\n")
+    exec("do_exec = exec")
+    def reraise(cls, val, tb):
+        raise val.with_traceback(tb)
     unicode = str
     _long_type = int
     from _thread import interrupt_main
@@ -1044,7 +1046,7 @@ def trace(msg):
 
     def executetask(self, item):
         try:
-            channel, (source, call_name, kwargs) = item
+            channel, (source, file_name, call_name, kwargs) = item
             if not ISPY3 and kwargs:
                 # some python2 versions do not accept unicode keyword params
                 # note: Unserializer generally turns py2-str to py3-str objects
@@ -1054,12 +1056,15 @@ def executetask(self, item):
                         name = name.encode('ascii')
                     newkwargs[name] = value
                 kwargs = newkwargs
+            if source is None:
+                assert file_name, file_name
+                source = "".join(linecache.updatecache(file_name))
             loc = {'channel': channel, '__name__': '__channelexec__'}
             self._trace("execution starts[%s]: %s" %
                         (channel.id, repr(source)[:50]))
             channel._executing = True
             try:
-                co = compile(source+'\n', '<remote exec>', 'exec')
+                co = compile(source+'\n', file_name or '<remote exec>', 'exec')
                 do_exec(co, loc)  # noqa
                 if call_name:
                     self._trace('calling %s(**%60r)' % (call_name, kwargs))

execnet/script/socketserver.py

@@ -63,7 +63,7 @@ def exec_from_one_connection(serversock):
     }
     source = eval(source)
     if source:
-        co = compile(source+'\n', source, 'exec')
+        co = compile(source+'\n', '<socket server>', 'exec')
         print_(progname, 'compiled source, executing')
         try:
             exec_(co, g)  # noqa

testing/test_basics.py

@@ -326,7 +326,7 @@ def test_source_of_nested_function(self):
         def working(channel):
             pass
 
-        send_source = gateway._source_of_function(working)
+        send_source = gateway._source_of_function(working).lstrip("\r\n")
         expected = 'def working(channel):\n    pass\n'
         assert send_source == expected
 

testing/test_gateway.py

@@ -110,6 +110,36 @@ def test_remote_exec_module(self, tmpdir, gw):
         name = channel.receive()
         assert name == 2
 
+    def test_remote_exec_module_with_traceback(self, gw, tmpdir):
+        remotetest = tmpdir.join("remotetest.py")
+        remotetest.write(
+            "\ndef run_me(channel=None):\n  raise ValueError('me')\n\n" +
+            "if __name__ == '__channelexec__':\n  run_me()\n")
+
+        try:
+            sys.path.insert(0, str(tmpdir))
+            module = __import__("remotetest")
+        finally:
+            sys.path.pop(0)
+
+        ch = gw.remote_exec(module)
+        try:
+            ch.receive()
+        except execnet.gateway_base.RemoteError as e:
+            assert 'remotetest.py", line 3, in run_me' in str(e)
+            assert "ValueError: me" in str(e)
+        finally:
+            ch.close()
+
+        ch = gw.remote_exec(module.run_me)
+        try:
+            ch.receive()
+        except execnet.gateway_base.RemoteError as e:
+            assert 'remotetest.py", line 3, in run_me' in str(e)
+            assert "ValueError: me" in str(e)
+        finally:
+            ch.close()
+
     def test_correct_setup_no_py(self, gw):
         channel = gw.remote_exec("""
             import sys