捕获并输出协程中的异常

wsrepl/framework.py

@@ -55,20 +55,22 @@ class Task:
         gen_log.debug('Task[%s] __init__ ', self.coro_id)
 
     def step(self, result=None):
-        try:
-            with self.ws_context():
+        future_or_none = None
+        with self.ws_context():
+            try:
                 future_or_none = self.coro.send(result)
-            if not isinstance(future_or_none, WebIOFuture) and future_or_none is not None:
-                if not self.ws.closed():
-                    future_or_none.add_done_callback(self._tornado_future_callback)
-                    self.pending_futures[id(future_or_none)] = future_or_none
-        except StopIteration as e:
-            if len(e.args) == 1:
-                self.result = e.args[0]
-
-            self.task_finished = True
-
-            gen_log.debug('Task[%s] finished', self.coro_id)
+            except StopIteration as e:
+                if len(e.args) == 1:
+                    self.result = e.args[0]
+                self.task_finished = True
+                gen_log.debug('Task[%s] finished', self.coro_id)
+            except Exception as e:
+                self.ws.on_coro_error()
+
+        if not isinstance(future_or_none, WebIOFuture) and future_or_none is not None:
+            if not self.ws.closed():
+                future_or_none.add_done_callback(self._tornado_future_callback)
+                self.pending_futures[id(future_or_none)] = future_or_none
 
     def _tornado_future_callback(self, future):
         del self.pending_futures[id(future)]

wsrepl/output.py

@@ -27,7 +27,7 @@ def json_print(obj):
 put_markdown = text_print
 
 
-def put_table(tdata):
+def put_table(tdata, header=None):
     """
     |      \|      |      |      |
     | ---- | ---- | ---- | ---- |
@@ -35,8 +35,14 @@ def put_table(tdata):
     |      |      |      |      |
     |      |      |      |      |
     :param tdata:
+    :param header: 列表，当tdata为字典列表时，header指定表头顺序
     :return:
     """
+    if header:
+        tdata = [
+            [row.get(k, '') for k in header]
+            for row in tdata
+        ]
 
     def quote(data):
         return data.replace('|', r'\|')
@@ -107,4 +113,4 @@ def put_file(name, content):
     :return:
     """
     content = b64encode(content).decode('ascii')
-    send_msg('output', dict(type='file', name=name, content=content))
+    send_msg('output', dict(type='file', name=name, content=content))

wsrepl/platform/tornado.py

@@ -9,11 +9,13 @@ from tornado.log import gen_log
 from wsrepl.framework import Task
 
 from wsrepl import project_dir
+import sys, traceback
+from wsrepl.output import put_markdown
 
 STATIC_PATH = '%s/html' % project_dir
 
 
-def ws_handler(coro_func):
+def ws_handler(coro_func, debug=True):
     class WSHandler(tornado.websocket.WebSocketHandler):
 
         def check_origin(self, origin):
@@ -71,6 +73,13 @@ def ws_handler(coro_func):
 
             self.step_task(coro, data)
 
+        def on_coro_error(self):
+            type, value, tb = sys.exc_info()
+            tb_len = len(list(traceback.walk_tb(tb)))
+            lines = traceback.format_exception(type, value, tb, limit=1 - tb_len)
+            traceback_msg = ''.join(lines)
+            put_markdown("发生错误：\n```\n%s\n```" % traceback_msg)
+
         def on_close(self):
             self._closed = True
             print("WebSocket closed")