From e0f18baf65f129ac65e71066beba6142e31139fe Mon Sep 17 00:00:00 2001
From: Anirban Basu <anirbanbasu@users.noreply.github.com>
Date: Sun, 8 Sep 2024 14:05:01 +0900
Subject: [PATCH] fix: Added a possible fix to get_terminal_size for
 Huggingface Spaces.

---
 src/dqa.py   |  8 +++++---
 src/utils.py | 22 ++++++++++++++++++++++
 2 files changed, 27 insertions(+), 3 deletions(-)

diff --git a/src/dqa.py b/src/dqa.py
index 125ffd2..813d707 100644
--- a/src/dqa.py
+++ b/src/dqa.py
@@ -20,7 +20,6 @@
     ic = lambda *a: None if not a else (a[0] if len(a) == 1 else a)  # noqa
 
 from tqdm import tqdm
-import os
 import asyncio
 import uuid
 
@@ -54,6 +53,7 @@
     EMPTY_STRING,
     FAKE_STRING,
     ToolNames,
+    get_terminal_size,
 )  # , parse_env, EnvironmentVariables
 
 
@@ -933,7 +933,7 @@ async def run(self, query: str):
         done: bool = False
         total_steps: int = 0
         finished_steps: int = 0
-        terminal_columns, _ = os.get_terminal_size()
+        terminal_columns, _ = get_terminal_size()
         progress_bar = tqdm(
             total=total_steps,
             leave=True,
@@ -945,7 +945,9 @@ async def run(self, query: str):
         async for ev in self.workflow.stream_events():
             total_steps = ev.total_steps
             finished_steps = ev.finished_steps
-            print(f"\n{str(ev.msg)}")
+            print(f"\n{str(ev.msg)}", flush=True)
+            # TODO: Is tqdm.write better than printf?
+            # tqdm.write(f"\n{str(ev.msg)}")
             progress_bar.reset(total=total_steps)
             progress_bar.update(finished_steps)
             progress_bar.refresh()
diff --git a/src/utils.py b/src/utils.py
index 751c140..c0f11b6 100644
--- a/src/utils.py
+++ b/src/utils.py
@@ -141,6 +141,28 @@ def parse_env(
     return value
 
 
+def get_terminal_size(fallback=(100, 25)) -> tuple[int, int]:
+    """
+    Get the terminal size.
+    See: https://granitosaur.us/getting-terminal-size.html
+
+    Args:
+        fallback (tuple[int, int]): The fallback size to use if the terminal size cannot be determined.
+
+    Returns:
+        tuple[int, int]: The terminal size as a tuple of (columns, rows).
+    """
+    for i in range(0, 3):
+        try:
+            columns, rows = os.get_terminal_size(i)
+        except OSError:
+            continue
+        break
+    else:  # set default if the loop completes which means all failed
+        columns, rows = fallback
+    return columns, rows
+
+
 def check_list_subset(list_a: list[Any], list_b: list[Any]) -> list[Any]:
     """
     Check if the elements of list_a forms a set that is a subset of the set formed by the elements of list_a.