Fix colors in Jupyter Notebook and Google Colab (#494)

CHANGELOG.rst

@@ -3,6 +3,7 @@
 
 - Add a new ``watch`` optional argument to file sinks in order to automatically re-create possibly deleted or changed file (`#471 <https://github.com/Delgan/loguru/issues/471>`_).
 - Make ``patch()`` calls cumulative instead of overriding the possibly existing patching function (`#462 <https://github.com/Delgan/loguru/issues/462>`_).
+- Fix logs colorization not automatically enabled for Jupyter Notebook and Google Colab (`#494 <https://github.com/Delgan/loguru/issues/494>`_).
 
 
 `0.6.0`_ (2022-01-29)

loguru/_colorama.py

@@ -6,6 +6,19 @@ def should_colorize(stream):
     if stream is None:
         return False
 
+    try:
+        import ipykernel
+        import IPython
+
+        ipython = IPython.get_ipython()
+        is_jupyter_stream = isinstance(stream, ipykernel.iostream.OutStream)
+        is_jupyter_shell = isinstance(ipython, ipykernel.zmqshell.ZMQInteractiveShell)
+    except Exception:
+        pass
+    else:
+        if is_jupyter_stream and is_jupyter_shell:
+            return True
+
     if stream is sys.__stdout__ or stream is sys.__stderr__:
         if "PYCHARM_HOSTED" in os.environ:
             return True

tests/test_add_option_colorize.py

@@ -2,9 +2,8 @@
 import sys
 from unittest.mock import MagicMock
 
-import pytest
-
 import loguru
+import pytest
 from loguru import logger
 
 from .conftest import parse
@@ -151,6 +150,53 @@ def test_pycharm_ignored(monkeypatch, stream):
     assert not loguru._colorama.should_colorize(stream)
 
 
+def test_jupyter_fixed(monkeypatch):
+    class Shell:
+        pass
+
+    class Out:
+        pass
+
+    stream = MagicMock()
+    stream.__class__ = Out
+    stream.isatty.return_value = False
+    ipython = MagicMock()
+    ipykernel = MagicMock()
+    instance = MagicMock()
+    instance.__class__ = Shell
+    ipython.get_ipython.return_value = instance
+    ipykernel.zmqshell.ZMQInteractiveShell = Shell
+    ipykernel.iostream.OutStream = Out
+
+    monkeypatch.setitem(sys.modules, "IPython", ipython)
+    monkeypatch.setitem(sys.modules, "ipykernel", ipykernel)
+
+    assert not stream.isatty()
+    assert loguru._colorama.should_colorize(stream)
+
+
+@pytest.mark.parametrize("stream", [None, Stream(False), Stream(None)])
+def test_jupyter_ignored(monkeypatch, stream):
+    class Shell:
+        pass
+
+    class Out:
+        pass
+
+    ipython = MagicMock()
+    ipykernel = MagicMock()
+    instance = MagicMock()
+    instance.__class__ = Shell
+    ipython.get_ipython.return_value = instance
+    ipykernel.zmqshell.ZMQInteractiveShell = Shell
+    ipykernel.iostream.OutStream = Out
+
+    monkeypatch.setitem(sys.modules, "IPython", ipython)
+    monkeypatch.setitem(sys.modules, "ipykernel", ipykernel)
+
+    assert not loguru._colorama.should_colorize(stream)
+
+
 @pytest.mark.parametrize("stream", [sys.__stdout__, sys.__stderr__])
 def test_mintty_fixed_windows(monkeypatch, stream):
     monkeypatch.setattr(os, "name", "nt")