Fix logging on colab.

riga · Jan 9, 2023 · b8d00e3 · b8d00e3
1 parent e92fb58
commit b8d00e3
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 1 deletion.
diff --git a/law/logger.py b/law/logger.py
@@ -15,7 +15,7 @@
 
 import six
 
-from law.util import no_value, colored, ipykernel
+from law.util import no_value, colored, ipykernel, ON_COLAB
 
 
 _logging_setup = False
@@ -34,6 +34,10 @@ def setup_logging():
         return
     _logging_setup = True
 
+    # remove root handlers on colab
+    if ON_COLAB and logging.root.handlers:
+        logging.root.removeHandler(logging.root.handlers[0])
+
     # setup the main law logger first and set its handler which is propagated to subloggers
     logger = get_logger("law", skip_setup=True)
     logger = setup_logger(logger, add_console_handler=False)

diff --git a/law/util.py b/law/util.py
@@ -54,6 +54,12 @@
 except ImportError:
     ipykernel = None
 
+try:
+    import google.colab  # noqa
+    ON_COLAB = True
+except ImportError:
+    ON_COLAB = False
+
 
 logger = logging.getLogger(__name__)