diff --git a/halo/halo.py b/halo/halo.py index 9e10b66..fecc883 100644 --- a/halo/halo.py +++ b/halo/halo.py @@ -35,11 +35,16 @@ class Halo(object): """ CLEAR_LINE = "\033[K" + CLEAR_REST = "\033[J" SPINNER_PLACEMENTS = ( "left", "right", ) + # a global list to keep all Halo instances + _instances = [] + _lock = threading.Lock() + def __init__( self, text="", @@ -96,6 +101,8 @@ def __init__( self._stop_spinner = None self._spinner_id = None self.enabled = enabled + self._stopped = False + self._content = "" environment = get_environment() @@ -294,7 +301,34 @@ def _check_stream(self): return True - def _write(self, s): + def _pop_stream_content_until_self(self, clear_self=False): + """Move cursor to the end of this instance's content and erase all contents + following it. + Parameters + ---------- + clear_self: bool + If equals True, the content of current line will also get cleared + Returns + ------- + str + The content of stream following this instance. + """ + erased_content = [] + lines_to_erase = self._content.count("\n") if clear_self else 0 + for inst in Halo._instances[::-1]: + if inst is self: + break + erased_content.append(inst._content) + lines_to_erase += inst._content.count("\n") + + if lines_to_erase > 0: + # Move cursor up n lines + self._write_stream("\033[{}A".format(lines_to_erase)) + # Erase rest content + self._write_stream(self.CLEAR_REST) + return "".join(reversed(erased_content)) + + def _write_stream(self, s): """Write to the stream, if writable Parameters ---------- @@ -304,15 +338,29 @@ def _write(self, s): if self._check_stream(): self._stream.write(s) - def _hide_cursor(self): - """Disable the user's blinking cursor + def _write(self, s, overwrite=False): + """Write to the stream and keep following lines unchanged. + Parameters + ---------- + s : str + Characters to write to the stream + overwrite: bool + If set to True, overwrite the content of current instance. """ + with Halo._lock: + erased_content = self._pop_stream_content_until_self(overwrite) + self._write_stream(s) + # Write back following lines + self._write_stream(erased_content) + self._content = s if overwrite else self._content + s + + def _hide_cursor(self): + """Disable the user's blinking cursor""" if self._check_stream() and self._stream.isatty(): cursor.hide(stream=self._stream) def _show_cursor(self): - """Re-enable the user's blinking cursor - """ + """Re-enable the user's blinking cursor""" if self._check_stream() and self._stream.isatty(): cursor.show(stream=self._stream) @@ -390,26 +438,26 @@ def clear(self): ------- self """ - self._write("\r") - self._write(self.CLEAR_LINE) + with Halo._lock: + erased_content = self._pop_stream_content_until_self(True) + self._content = "" + self._write_stream(erased_content) return self def _render_frame(self): - """Renders the frame on the line after clearing it. - """ + """Renders the frame on the line after clearing it.""" if not self.enabled: # in case we're disabled or stream is closed while still rendering, # we render the frame and increment the frame index, so the proper # frame is rendered if we're reenabled or the stream opens again. return - self.clear() frame = self.frame() - output = "\r{}".format(frame) + output = "\r{}\n".format(frame) try: - self._write(output) + self._write(output, True) except UnicodeEncodeError: - self._write(encode_utf_8_text(output)) + self._write(encode_utf_8_text(output), True) def render(self): """Runs the render until thread flag is set. @@ -490,6 +538,14 @@ def start(self, text=None): if not (self.enabled and self._check_stream()): return self + # Clear all stale Halo instances created before + # Check against Halo._instances instead of self._instances + # to avoid possible overriding in subclasses. + if all(inst._stopped for inst in Halo._instances): + Halo._instances[:] = [] + # Allow for calling start() multiple times + if self not in Halo._instances: + Halo._instances.append(self) self._hide_cursor() self._stop_spinner = threading.Event() @@ -498,6 +554,7 @@ def start(self, text=None): self._render_frame() self._spinner_id = self._spinner_thread.name self._spinner_thread.start() + self._stopped = False return self @@ -511,12 +568,17 @@ def stop(self): self._stop_spinner.set() self._spinner_thread.join() + if self._stopped: + return + if self.enabled: self.clear() self._frame_index = 0 self._spinner_id = None self._show_cursor() + self._stopped = True + return self def succeed(self, text=None): diff --git a/tests/test_halo.py b/tests/test_halo.py index 29918aa..17d5785 100644 --- a/tests/test_halo.py +++ b/tests/test_halo.py @@ -622,8 +622,26 @@ def test_redirect_stdout(self): self.assertIn('foo', output[0]) + def test_running_multiple_instances(self): + spinner = Halo("foo", stream=self._stream) + _instance = Halo() + # Pretend that another spinner is being displayed under spinner + _instance._content = "Some lines\n" + Halo._instances.extend([spinner, _instance]) + spinner.start() + time.sleep(1) + spinner.stop() + spinner.stop_and_persist(text="Done") + output = self._get_test_output()["text"] + self.assertEqual(output[0], "{} foo".format(frames[0])) + self.assertEqual(output[1], "Some lines") + self.assertEqual(output[2], "{} foo".format(frames[1])) + self.assertEqual(output[3], "Some lines") + self.assertEqual(output[-2], " Done") + self.assertEqual(output[-1], "Some lines") + def tearDown(self): - pass + self._stream.close() if __name__ == '__main__':