nit

Signed-off-by: Kevin Su <pingsutw@apache.org>
flyteorg · Jul 15, 2024 · c1b704c · c1b704c
1 parent 165ecba
commit c1b704c
Showing 1 changed file with 6 additions and 5 deletions.
diff --git a/flytekit/clis/sdk_in_container/utils.py b/flytekit/clis/sdk_in_container/utils.py
@@ -133,11 +133,12 @@ def pretty_print_traceback(e: Exception, verbosity: int = 1):
 
     if isinstance(e, FlyteCompilationException):
         e = annotate_exception_with_code(e, e.fn, e.param_name)
-        # TODO: Use other way to check if the background is light or dark
-        theme = "emacs" if "LIGHT_BACKGROUND" in os.environ else "monokai"
-        syntax = Syntax(getattr(e, SOURCE_CODE), "python", theme=theme, background_color="default")
-        panel = Panel(syntax, border_style="red", title=e._ERROR_CODE, title_align="left")
-        console.print(panel, no_wrap=False)
+        if hasattr(e, SOURCE_CODE):
+            # TODO: Use other way to check if the background is light or dark
+            theme = "emacs" if "LIGHT_BACKGROUND" in os.environ else "monokai"
+            syntax = Syntax(getattr(e, SOURCE_CODE), "python", theme=theme, background_color="default")
+            panel = Panel(syntax, border_style="red", title=e._ERROR_CODE, title_align="left")
+            console.print(panel, no_wrap=False)
 
 
 def pretty_print_exception(e: Exception, verbosity: int = 1):