Simplify traceback when calling Ploomber task - closes #605

src/ploomber/cli/io.py

@@ -1,11 +1,11 @@
 from functools import wraps
 import sys
-import traceback
 
 import click
 
 from ploomber.io import TerminalWriter
 from ploomber.exceptions import DAGBuildError, DAGRenderError, BaseException
+from ploomber.executors import _format
 
 # TODO: there are two types of cli commands: the ones that execute user's
 # code (ploomber build/task) and the ones that parse a dag/task but do not
@@ -42,8 +42,8 @@ def wrapper(catch_exception=True, **kwargs):
             except BaseException as e:
                 click.secho(e.get_message(), file=sys.stderr, fg='red')
                 sys.exit(1)
-            except Exception:
-                error = traceback.format_exc()
+            except Exception as e:
+                error = _format.exception(e)
                 color = True
             else:
                 error = None

src/ploomber/executors/_format.py

@@ -1,6 +1,7 @@
 import traceback
 
 from ploomber.exceptions import TaskBuildError, RenderError, TaskRenderError
+
 from papermill.exceptions import PapermillExecutionError
 
 

tests/cli/test_custom.py

@@ -488,3 +488,27 @@ def test_build_with_replaced_env_value(tmp_nbs, monkeypatch_session):
         sys, 'argv',
         ['python', '--entry-point', 'pipeline.yaml', '--env--sample', 'True'])
     build.main(catch_exception=False)
+
+
+def test_task_command_formats_exception(tmp_directory, monkeypatch,
+                                        tmp_imports, capsys):
+    Path('functions.py').write_text("""
+def load(product):
+    raise ValueError
+""")
+
+    Path('pipeline.yaml').write_text("""
+tasks:
+    - source: functions.load
+      product: output.txt
+""")
+
+    monkeypatch.setattr(sys, 'argv', ['ploomber', 'task', 'load'])
+
+    with pytest.raises(SystemExit):
+        cmd_router()
+
+    captured = capsys.readouterr()
+    expected = ('raise ValueError\nValueError\n\nploomber.exceptions.'
+                'TaskBuildError: Error building task "load"')
+    assert expected in captured.err