From ae808894771a09ca5c8f835258009b7e49e2b574 Mon Sep 17 00:00:00 2001
From: skshetry <18718008+skshetry@users.noreply.github.com>
Date: Tue, 7 Feb 2023 10:05:05 +0545
Subject: [PATCH] logging: defer calling colorama.init() in jupyter, otherwise
 initialize immediately (#8969)Closes
 https://github.com/iterative/dvc/issues/8957

logging: init colorama before logging if outside jupyter env

Otherwise, tty detection does not work as `colorama.init()` mutates
sys.stdout/sys.stderr, and after `dictConfig()`, the handlers will hold
on to the stream.
---
 dvc/logger.py | 24 ++++++++++++++++++++----
 1 file changed, 20 insertions(+), 4 deletions(-)

diff --git a/dvc/logger.py b/dvc/logger.py
index 21c943653..7aa294a98 100644
--- a/dvc/logger.py
+++ b/dvc/logger.py
@@ -203,7 +203,22 @@ def set_loggers_level(level: int = logging.INFO) -> None:
         logging.getLogger(name).setLevel(level)
 
 
+def _is_jupyter() -> bool:
+    try:
+        ipython = get_ipython()  # type: ignore[name-defined]
+    except (NameError, TypeError):
+        return False
+
+    shell = ipython.__class__.__name__
+    return shell == "ZMQInteractiveShell" or "google.colab" in str(ipython.__class__)
+
+
 def setup(level: int = logging.INFO) -> None:
+    is_jupyter = _is_jupyter()
+    if not is_jupyter:
+        # initialize colorama, so that tty detection works if we are not inside jupyter
+        colorama.init()
+
     if level >= logging.DEBUG:
         # Unclosed session errors for asyncio/aiohttp are only available
         # on the tracing mode for extensive debug purposes. They are really
@@ -287,7 +302,8 @@ def setup(level: int = logging.INFO) -> None:
         }
     )
 
-    # NOTE: colorama init needs to go after logging config to avoid potential
-    # conflicts with libs that modify logging (like tensorflow/abseil)
-    # https://github.com/iterative/dvc/issues/8387
-    colorama.init()
+    if is_jupyter:
+        # NOTE: colorama init needs to go after logging config to avoid potential
+        # conflicts with libs that modify logging (like tensorflow/abseil)
+        # https://github.com/iterative/dvc/issues/8387
+        colorama.init()