Use alternative to ANSI control sequences in Jupyter notebooks (#195)

Second try of #193 and fixing #176 - Checking if stdout is a TTY to determine if output is a Jupyter cell or a terminal - Disabled ANSI control sequence DECTCEM for Jupyter (anyway there is no cursor) - Replaced ANSI control sequence EL for Jupyter (where it doesn't erase the line, but just prints three white spaces) with some code that simply overrides the entire line with white spaces - Added patching sys.stdout.isatty to pretend output is a terminal
pavdmyt · Jun 3, 2022 · 6e96e83 · 6e96e83
1 parent de08a47
commit 6e96e83
Show file tree

Hide file tree

Showing 5 changed files with 120 additions and 22 deletions.
diff --git a/Makefile b/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)"

diff --git a/tests/conftest.py b/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

diff --git a/tests/test_attrs.py b/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]

diff --git a/tests/test_in_out.py b/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)
diff --git a/yaspin/core.py b/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))