From b7f6b746d37401d903d32f3df2e8594263ff8d68 Mon Sep 17 00:00:00 2001
From: Brian Pugh <bnp117@gmail.com>
Date: Mon, 24 Apr 2023 08:34:46 -0700
Subject: [PATCH] Fully remove belay-internal stacktrace for install/run/exec
 cli commands

---
 belay/cli/common.py  | 14 ++++++++++++++
 belay/cli/exec.py    |  5 +++--
 belay/cli/install.py |  5 +++--
 belay/cli/run.py     |  4 ++--
 4 files changed, 22 insertions(+), 6 deletions(-)

diff --git a/belay/cli/common.py b/belay/cli/common.py
index 259f6af..66d7fd8 100644
--- a/belay/cli/common.py
+++ b/belay/cli/common.py
@@ -1,2 +1,16 @@
+from belay.pyboard import PyboardException
+
 help_port = "Port (like /dev/ttyUSB0) or WebSocket (like ws://192.168.1.100) of device."
 help_password = "Password for communication methods (like WebREPL) that require authentication."  # nosec  # noqa: S105
+
+
+class remove_stacktrace:  # noqa: N801
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        if exc_type is not None and issubclass(exc_type, PyboardException):
+            print(exc_value)
+            return True  # suppress the full stack trace
+        else:
+            return False  # let other exceptions propagate normally
diff --git a/belay/cli/exec.py b/belay/cli/exec.py
index e155f7a..04bc6f5 100644
--- a/belay/cli/exec.py
+++ b/belay/cli/exec.py
@@ -3,7 +3,7 @@
 from typer import Argument, Option
 
 from belay import Device
-from belay.cli.common import help_password, help_port
+from belay.cli.common import help_password, help_port, remove_stacktrace
 
 
 def exec(
@@ -13,5 +13,6 @@ def exec(
 ):
     """Execute python statement on-device."""
     device = Device(port, password=password)
-    device(statement)
+    with remove_stacktrace():
+        device(statement)
     device.close()
diff --git a/belay/cli/install.py b/belay/cli/install.py
index fd9cd99..09e78f9 100644
--- a/belay/cli/install.py
+++ b/belay/cli/install.py
@@ -8,7 +8,7 @@
 from typer import Argument, Option
 
 from belay import Device
-from belay.cli.common import help_password, help_port
+from belay.cli.common import help_password, help_port, remove_stacktrace
 from belay.cli.sync import sync_device as _sync_device
 from belay.project import find_project_folder, load_groups, load_pyproject
 
@@ -97,7 +97,8 @@ def progress_update(description=None, **kwargs):
 
         if run:
             content = run.read_text()
-            device(content)
+            with remove_stacktrace():
+                device(content)
             return
 
         # Reset device so ``main.py`` has a chance to execute.
diff --git a/belay/cli/run.py b/belay/cli/run.py
index f421459..a7cc527 100644
--- a/belay/cli/run.py
+++ b/belay/cli/run.py
@@ -3,7 +3,7 @@
 from typer import Argument, Option
 
 from belay import Device
-from belay.cli.common import help_password, help_port
+from belay.cli.common import help_password, help_port, remove_stacktrace
 
 
 def run(
@@ -23,5 +23,5 @@ def run(
 
     """
     content = file.read_text()
-    with Device(port, password=password) as device:
+    with Device(port, password=password) as device, remove_stacktrace():
         device(content)