Use alternative to ANSI control sequences in Jupyter notebooks

Makefile

@@ -13,7 +13,7 @@ pypi_pwd := $(shell grep password ~/.pypirc | awk -F"= " '{ print $$2 }')
 
 flake:
 	@echo "$(OK_COLOR)==> Linting code ...$(NO_COLOR)"
-	@poetry run flake8 --ignore=F821,E501 .
+	@poetry run flake8 --ignore=F821,E501,W503 .
 
 lint:
 	@echo "$(OK_COLOR)==> Linting code ...$(NO_COLOR)"

tests/conftest.py

@@ -95,6 +95,16 @@ def reversal(request):
     return request.param
 
 
+@pytest.fixture(scope="session", params=[True, False], ids=["terminal", "jupyter"])
+def isatty_fixture(request):
+    return request.param
+
+
+@pytest.fixture(autouse=True)
+def isatty_true(monkeypatch):
+    monkeypatch.setattr(sys.stdout, "isatty", lambda: True)
+
+
 def color_id_func(case):
     if isinstance(case, tuple):
         color, _ = case

tests/test_attrs.py

@@ -5,6 +5,8 @@
 Test Yaspin attributes magic hidden in __getattr__.
 """
 
+import sys
+
 import pytest
 
 from yaspin import yaspin
@@ -19,12 +21,13 @@ def test_set_spinner_by_name(attr_name):
 
 
 # Values for ``color`` argument
-def test_color(color_test_cases):
+def test_color(monkeypatch, color_test_cases):
     color, expected = color_test_cases
     # ``None`` and ``""`` are skipped
     if not color:
         pytest.skip("{0} - unsupported case".format(repr(color)))
 
+    monkeypatch.setattr(sys.stdout, "isatty", lambda: True)
     sp = yaspin()
 
     if isinstance(expected, Exception):
@@ -37,12 +40,13 @@ def test_color(color_test_cases):
 
 
 # Values for ``on_color`` argument
-def test_on_color(on_color_test_cases):
+def test_on_color(monkeypatch, on_color_test_cases):
     on_color, expected = on_color_test_cases
     # ``None`` and ``""`` are skipped
     if not on_color:
         pytest.skip("{0} - unsupported case".format(repr(on_color)))
 
+    monkeypatch.setattr(sys.stdout, "isatty", lambda: True)
     sp = yaspin()
 
     if isinstance(expected, Exception):
@@ -60,7 +64,8 @@ def test_on_color(on_color_test_cases):
 @pytest.mark.parametrize(
     "attr", sorted([k for k, v in COLOR_MAP.items() if v == "attrs"])
 )
-def test_attrs(attr):
+def test_attrs(monkeypatch, attr):
+    monkeypatch.setattr(sys.stdout, "isatty", lambda: True)
     sp = yaspin()
     getattr(sp, attr)
     assert sp.attrs == [attr]

tests/test_in_out.py

@@ -75,22 +75,48 @@ def test_compose_out_with_color(
     assert isinstance(out, str)
 
 
-def test_write(capsys, text):
+def test_color_jupyter(monkeypatch):
+    monkeypatch.setattr(sys.stdout, "isatty", lambda: False)
+    with pytest.warns(UserWarning):
+        sp = yaspin(color="red")
+
+    out = sp._compose_out(frame=u"/")
+    assert "\033" not in out
+
+
+def test_write(monkeypatch, capsys, text, isatty_fixture):
+    monkeypatch.setattr(sys.stdout, "isatty", lambda: isatty_fixture)
     sp = yaspin()
     sp.write(text)
 
     out, _ = capsys.readouterr()
-    # cleans stdout from _clear_line and \r
-    out = out.replace("\r\033[0K", "")
+    # cleans stdout from _clear_line
+    if isatty_fixture:
+        out = out.replace("\r\033[K", "")
+    else:
+        out = out.replace("\r", "")
 
     assert isinstance(out, (str, bytes))
     assert out[-1] == "\n"
     if text:
         assert out[:-1] == text
 
 
-def test_hide_show(capsys, text, request):
+def test_show_jupyter(monkeypatch, capsys):
+    monkeypatch.setattr(sys.stdout, "isatty", lambda: False)
+    with yaspin(text="12345") as sp:
+        sp.start()
+        sp.write("123")
+
+    out, _ = capsys.readouterr()
+    # check spinner line was correctly overridden with whitespaces
+    # r = \r, s = spinner char, w = space, 12345 = printed chars
+    assert "12345\r" + " " * len("rsw12345") + "\r123" in out
+
+
+def test_hide_show(monkeypatch, capsys, text, request, isatty_fixture):
     # Setup
+    monkeypatch.setattr(sys.stdout, "isatty", lambda: isatty_fixture)
     sp = yaspin()
     sp.start()
 
@@ -110,15 +136,21 @@ def teardown():
     out, _ = capsys.readouterr()
 
     # ensure that text was cleared with the hide method
-    assert out[-5:] == "\r\033[0K"
+    if isatty_fixture:
+        assert out[-4:] == "\r\033[K"
+    else:
+        assert out[-1:] == "\r"
 
     # ``\n`` is required to flush stdout during
     # the hidden state of the spinner
     sys.stdout.write("{0}\n".format(text))
     out, _ = capsys.readouterr()
 
-    # cleans stdout from _clear_line and \r
-    out = out.replace("\r\033[0K", "")
+    # cleans stdout from _clear_line
+    if isatty_fixture:
+        out = out.replace("\r\033[K", "")
+    else:
+        out = out.replace("\r", "")
 
     assert isinstance(out, (str, bytes))
     assert out[-1] == "\n"
@@ -132,7 +164,10 @@ def teardown():
     out, _ = capsys.readouterr()
 
     # ensure that text was cleared before resuming the spinner
-    assert out[:5] == "\r\033[0K"
+    if isatty_fixture:
+        assert out[:4] == "\r\033[K"
+    else:
+        assert out[:1] == "\r"
 
 
 def test_spinner_write_race_condition(capsys):
@@ -154,9 +189,10 @@ def test_spinner_write_race_condition(capsys):
     assert not re.search(r"aaaa[^\rb]*bbbb", out)
 
 
-def test_spinner_hiding_with_context_manager(capsys):
+def test_spinner_hiding_with_context_manager(monkeypatch, capsys, isatty_fixture):
     HIDDEN_START = "hidden start"
     HIDDEN_END = "hidden end"
+    monkeypatch.setattr(sys.stdout, "isatty", lambda: isatty_fixture)
     sp = yaspin(text="foo")
     sp.start()
 
@@ -174,13 +210,19 @@ def test_spinner_hiding_with_context_manager(capsys):
 
     # make sure no spinner text was printed while the spinner was hidden
     out, _ = capsys.readouterr()
-    out = out.replace("\r\033[0K", "")
+    if isatty_fixture:
+        out = out.replace("\r\033[K", "")
+    else:
+        out = out.replace("\r", "")
     assert "{}\n{}".format(HIDDEN_START, HIDDEN_END) in out
 
 
-def test_spinner_nested_hiding_with_context_manager(capsys):
+def test_spinner_nested_hiding_with_context_manager(
+    monkeypatch, capsys, isatty_fixture
+):
     HIDDEN_START = "hidden start"
     HIDDEN_END = "hidden end"
+    monkeypatch.setattr(sys.stdout, "isatty", lambda: isatty_fixture)
     sp = yaspin(text="foo")
     sp.start()
 
@@ -202,7 +244,10 @@ def test_spinner_nested_hiding_with_context_manager(capsys):
 
     # make sure no spinner text was printed while the spinner was hidden
     out, _ = capsys.readouterr()
-    out = out.replace("\r\033[0K", "")
+    if isatty_fixture:
+        out = out.replace("\r\033[K", "")
+    else:
+        out = out.replace("\r", "")
     assert "{}\n{}".format(HIDDEN_START, HIDDEN_END) in out
 
 
@@ -234,9 +279,13 @@ def test_spinner_hiding_with_context_manager_and_exception():
         (["foo", "bar", "'", 23], """['foo', 'bar', "'", 23]"""),
     ],
 )
-def test_write_non_str_objects(capsys, obj, obj_str):
+def test_write_non_str_objects(monkeypatch, capsys, obj, obj_str, isatty_fixture):
+    monkeypatch.setattr(sys.stdout, "isatty", lambda: isatty_fixture)
     sp = yaspin()
     capsys.readouterr()
     sp.write(obj)
     out, _ = capsys.readouterr()
-    assert out == "\r\033[0K{}\n".format(obj_str)
+    if isatty_fixture:
+        assert out == "\r\033[K{}\n".format(obj_str)
+    else:
+        assert out == "\r\r{}\n".format(obj_str)

yaspin/core.py

@@ -16,6 +16,7 @@
 import sys
 import threading
 import time
+import warnings
 from typing import List, Set, Union
 
 from termcolor import colored
@@ -75,6 +76,7 @@ def __init__(  # pylint: disable=too-many-arguments
         self._last_frame = None
         self._stdout_lock = threading.Lock()
         self._hidden_level = 0
+        self._cur_line_len = 0
 
         # Signals
 
@@ -308,6 +310,7 @@ def write(self, text):
             assert isinstance(_text, str)
 
             sys.stdout.write("{0}\n".format(_text))
+            self._cur_line_len = 0
 
     def ok(self, text="OK"):
         """Set Ok (success) finalizer to a spinner."""
@@ -322,6 +325,13 @@ def fail(self, text="FAIL"):
     #
     # Protected
     #
+    @staticmethod
+    def _warn_color_disabled():
+        warnings.warn(
+            "color, on_color and attrs are not supported when running in jupyter",
+            stacklevel=3,
+        )
+
     def _freeze(self, final_text):
         """Stop spinner, compose last frame and 'freeze' it."""
         text = to_unicode(final_text)
@@ -332,6 +342,7 @@ def _freeze(self, final_text):
         self.stop()
         with self._stdout_lock:
             sys.stdout.write(self._last_frame)
+            self._cur_line_len = 0
 
     def _spin(self):
         while not self._stop_spin.is_set():
@@ -350,11 +361,16 @@ def _spin(self):
                 self._clear_line()
                 sys.stdout.write(out)
                 sys.stdout.flush()
+                self._cur_line_len = max(self._cur_line_len, len(out))
 
             # Wait
             self._stop_spin.wait(self._interval)
 
     def _compose_color_func(self):
+        if self.is_jupyter():
+            # ANSI Color Control Sequences are problematic in Jupyter
+            return None
+
         return functools.partial(
             colored,
             color=self._color,
@@ -428,8 +444,15 @@ def _reset_signal_handlers(self):
     #
     # Static
     #
+    @staticmethod
+    def is_jupyter() -> bool:
+        return not sys.stdout.isatty()
+
     @staticmethod
     def _set_color(value: str) -> str:
+        if Yaspin.is_jupyter():
+            Yaspin._warn_color_disabled()
+
         available_values = [k for k, v in COLOR_MAP.items() if v == "color"]
         if value not in available_values:
             raise ValueError(
@@ -441,6 +464,9 @@ def _set_color(value: str) -> str:
 
     @staticmethod
     def _set_on_color(value: str) -> str:
+        if Yaspin.is_jupyter():
+            Yaspin._warn_color_disabled()
+
         available_values = [k for k, v in COLOR_MAP.items() if v == "on_color"]
         if value not in available_values:
             raise ValueError(
@@ -451,6 +477,9 @@ def _set_on_color(value: str) -> str:
 
     @staticmethod
     def _set_attrs(attrs: List[str]) -> Set[str]:
+        if Yaspin.is_jupyter():
+            Yaspin._warn_color_disabled()
+
         available_values = [k for k, v in COLOR_MAP.items() if v == "attrs"]
         for attr in attrs:
             if attr not in available_values:
@@ -524,16 +553,21 @@ def _set_cycle(frames):
     @staticmethod
     def _hide_cursor():
         if sys.stdout.isatty():
+            # ANSI Control Sequence DECTCEM 1 does not work in Jupyter
             sys.stdout.write("\033[?25l")
             sys.stdout.flush()
 
     @staticmethod
     def _show_cursor():
         if sys.stdout.isatty():
+            # ANSI Control Sequence DECTCEM 2 does not work in Jupyter
             sys.stdout.write("\033[?25h")
             sys.stdout.flush()
 
-    @staticmethod
-    def _clear_line():
-        sys.stdout.write("\r")
-        sys.stdout.write("\033[0K")
+    def _clear_line(self):
+        if sys.stdout.isatty():
+            # ANSI Control Sequence EL does not work in Jupyter
+            sys.stdout.write("\r\033[K")
+        else:
+            fill = " " * self._cur_line_len
+            sys.stdout.write("\r{0}\r".format(fill))