From 743ef90ebe2769f3b0f0394af11ac2b9a674a993 Mon Sep 17 00:00:00 2001
From: Kevin Su <pingsutw@apache.org>
Date: Thu, 4 Jul 2024 08:45:41 +0800
Subject: [PATCH] Set default width if fail to get terminal size (#2558)

* Set default width if fail to get terminal size

Signed-off-by: Kevin Su <pingsutw@apache.org>

* lint

Signed-off-by: Kevin Su <pingsutw@apache.org>

* test

Signed-off-by: Kevin Su <pingsutw@apache.org>

---------

Signed-off-by: Kevin Su <pingsutw@apache.org>
Signed-off-by: Jan Fiedler <jan@union.ai>
---
 flytekit/core/type_engine.py |  2 +-
 flytekit/loggers.py          |  8 +++++++-
 tests/flytekit/loggers.py    | 11 +++++++++++
 3 files changed, 19 insertions(+), 2 deletions(-)
 create mode 100644 tests/flytekit/loggers.py

diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py
index 2a087cb4ca..92a4373839 100644
--- a/flytekit/core/type_engine.py
+++ b/flytekit/core/type_engine.py
@@ -1650,7 +1650,7 @@ def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type:
                 res_tag = trans.name
                 found_res = True
             except Exception as e:
-                logger.debug(f"Failed to convert from {lv} to {v}", e)
+                logger.debug(f"Failed to convert from {lv} to {v} with error: {e}")
 
         if is_ambiguous:
             raise TypeError(
diff --git a/flytekit/loggers.py b/flytekit/loggers.py
index 12b3e90423..8e38c2b4ee 100644
--- a/flytekit/loggers.py
+++ b/flytekit/loggers.py
@@ -140,13 +140,19 @@ def upgrade_to_rich_logging(log_level: typing.Optional[int] = logging.WARNING):
 
     import flytekit
 
+    try:
+        width = os.get_terminal_size().columns
+    except Exception as e:
+        logger.debug(f"Failed to get terminal size: {e}")
+        width = 80
+
     handler = RichHandler(
         tracebacks_suppress=[click, flytekit],
         rich_tracebacks=True,
         omit_repeated_times=False,
         show_path=False,
         log_time_format="%H:%M:%S.%f",
-        console=Console(width=os.get_terminal_size().columns),
+        console=Console(width=width),
     )
 
     formatter = logging.Formatter(fmt="%(filename)s:%(lineno)d - %(message)s")
diff --git a/tests/flytekit/loggers.py b/tests/flytekit/loggers.py
new file mode 100644
index 0000000000..3b94259049
--- /dev/null
+++ b/tests/flytekit/loggers.py
@@ -0,0 +1,11 @@
+import sys
+from subprocess import run
+
+
+def test_error(tmp_path):
+    EXAMPLE = "import flytekit"
+
+    script_path = tmp_path / "script.py"
+    script_path.write_text(EXAMPLE)
+    out = run([sys.executable, script_path], text=True, capture_output=True)
+    assert out.stderr == ""