diff --git a/click_spinner/__init__.py b/click_spinner/__init__.py index 3152d44..47bb309 100644 --- a/click_spinner/__init__.py +++ b/click_spinner/__init__.py @@ -13,11 +13,12 @@ def __init__(self, beep=False, disable=False, force=False, stream=sys.stdout): self.stream = stream self.stop_running = None self.spin_thread = None + self.tty_output = self.stream.isatty() or self.force def start(self): if self.disable: return - if self.stream.isatty() or self.force: + if self.tty_output: self.stop_running = threading.Event() self.spin_thread = threading.Thread(target=self.init_spin) self.spin_thread.start() @@ -43,8 +44,11 @@ def __exit__(self, exc_type, exc_val, exc_tb): if self.disable: return False self.stop() - if self.beep: - self.stream.write('\7') + if self.tty_output: + if self.beep: + self.stream.write('\7') + self.stream.flush() + self.stream.write(' \b') self.stream.flush() return False @@ -64,6 +68,8 @@ def spinner(beep=False, disable=False, force=False, stream=sys.stdout): Hide spinner. force : bool Force creation of spinner even when stdout is redirected. + stream : IO + Stream to write the spinner to. Example ------- diff --git a/tests/test_spinner.py b/tests/test_spinner.py index 3f6d1ed..cbdb091 100644 --- a/tests/test_spinner.py +++ b/tests/test_spinner.py @@ -58,6 +58,34 @@ def cli(): result = runner.invoke(cli, []) assert result.exception is None +def test_spinner_redirect_with_beep(): + @click.command() + def cli(): + stdout_io = StringIO() + with click_spinner.Spinner(beep=True, stream=stdout_io): + time.sleep(1) # allow time for a few spins + stdout_str = stdout_io.getvalue() + assert len(stdout_str) == 0 + + runner = CliRunner() + result = runner.invoke(cli, []) + assert result.exception is None + + +def test_spinner_erase(): + @click.command() + def cli(): + stdout_io = StringIO() + stdout_io.isatty = lambda: True + with click_spinner.Spinner(stream=stdout_io): + time.sleep(1) # allow time for a few spins + stdout_str = stdout_io.getvalue() + assert stdout_str[-3:] == '\b \b' + + runner = CliRunner() + result = runner.invoke(cli, []) + assert result.exception is None + def test_spinner_redirect_force(): @click.command()