From 70fb171670a3ec5124b2c01cae627ae0249ba845 Mon Sep 17 00:00:00 2001
From: ara-25 <abdurrehmanali25@gmail.com>
Date: Wed, 23 Oct 2024 15:16:35 +0500
Subject: [PATCH] Add traceback display with rich

---
 pyproject.toml                 |  3 ++-
 requirements-dev.txt           |  8 ++++++++
 requirements-docs.txt          | 10 +++++++++-
 src/mlinfra/cli/__init__.py    |  6 ++++++
 src/mlinfra/utils/constants.py |  2 ++
 src/mlinfra/utils/utils.py     | 10 ++++++++++
 6 files changed, 37 insertions(+), 2 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 14fc9db3..0d21af60 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -34,7 +34,8 @@ dependencies = [
     "GitPython",
     "mypy",
     "getmac",
-    "requests"
+    "requests",
+    "rich"
 ]
 
 [tool.setuptools]
diff --git a/requirements-dev.txt b/requirements-dev.txt
index 1cb730c8..c832139d 100644
--- a/requirements-dev.txt
+++ b/requirements-dev.txt
@@ -34,6 +34,10 @@ jmespath==1.0.1
     # via
     #   boto3
     #   botocore
+markdown-it-py==3.0.0
+    # via rich
+mdurl==0.1.2
+    # via markdown-it-py
 mypy==1.12.0
     # via mlinfra (pyproject.toml)
 mypy-extensions==1.0.0
@@ -48,6 +52,8 @@ pluggy==1.5.0
     # via pytest
 pre-commit==3.7.1
     # via mlinfra (pyproject.toml)
+pygments==2.18.0
+    # via rich
 pytest==8.2.2
     # via
     #   mlinfra (pyproject.toml)
@@ -62,6 +68,8 @@ pyyaml==6.0.2
     #   pre-commit
 requests==2.32.3
     # via mlinfra (pyproject.toml)
+rich==13.9.3
+    # via mlinfra (pyproject.toml)
 s3transfer==0.10.3
     # via boto3
 six==1.16.0
diff --git a/requirements-docs.txt b/requirements-docs.txt
index 32d75a9c..48728f32 100644
--- a/requirements-docs.txt
+++ b/requirements-docs.txt
@@ -50,10 +50,14 @@ markdown==3.6
     #   mkdocs
     #   mkdocs-material
     #   pymdown-extensions
+markdown-it-py==3.0.0
+    # via rich
 markupsafe==2.1.5
     # via
     #   jinja2
     #   mkdocs
+mdurl==0.1.2
+    # via markdown-it-py
 mergedeep==1.3.4
     # via
     #   mkdocs
@@ -89,7 +93,9 @@ pathspec==0.12.1
 platformdirs==4.3.6
     # via mkdocs-get-deps
 pygments==2.18.0
-    # via mkdocs-material
+    # via
+    #   mkdocs-material
+    #   rich
 pymdown-extensions==10.8.1
     # via mkdocs-material
 pyparsing==3.2.0
@@ -118,6 +124,8 @@ requests==2.32.3
     # via
     #   mlinfra (pyproject.toml)
     #   mkdocs-material
+rich==13.9.3
+    # via mlinfra (pyproject.toml)
 s3transfer==0.10.3
     # via boto3
 six==1.16.0
diff --git a/src/mlinfra/cli/__init__.py b/src/mlinfra/cli/__init__.py
index 76cf67a3..094cd0ed 100644
--- a/src/mlinfra/cli/__init__.py
+++ b/src/mlinfra/cli/__init__.py
@@ -9,3 +9,9 @@
 #     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
 #     or implied. See the License for the specific language governing
 #     permissions and limitations under the License.
+
+from mlinfra.utils.constants import ENABLE_RICH_TRACEBACK
+from mlinfra.utils.utils import setup_rich_traceback
+
+if ENABLE_RICH_TRACEBACK:
+    setup_rich_traceback()
diff --git a/src/mlinfra/utils/constants.py b/src/mlinfra/utils/constants.py
index 75150a85..825fc538 100644
--- a/src/mlinfra/utils/constants.py
+++ b/src/mlinfra/utils/constants.py
@@ -29,3 +29,5 @@
 VERSION: Final = version("mlinfra")
 
 EXECUTABLE_PATH = "/usr/local/bin"
+
+ENABLE_RICH_TRACEBACK = True
diff --git a/src/mlinfra/utils/utils.py b/src/mlinfra/utils/utils.py
index cd6c4002..e1092acc 100644
--- a/src/mlinfra/utils/utils.py
+++ b/src/mlinfra/utils/utils.py
@@ -21,6 +21,9 @@
 import sys
 from typing import List
 
+import invoke
+from rich.traceback import install as rich_tr_install
+
 from .constants import TF_PATH
 
 
@@ -186,3 +189,10 @@ def func_wrapper(*args, **kwargs):  # type: ignore
                 return None
 
     return func_wrapper
+
+
+def setup_rich_traceback(debug_mode: bool = False):
+    if debug_mode:
+        rich_tr_install(show_locals=True)
+    else:
+        rich_tr_install(show_locals=False, suppress=[invoke, subprocess])