From a2533e9a46f76c9d0e77fc12e5913807b30ae98f Mon Sep 17 00:00:00 2001 From: Kovid Goyal Date: Thu, 11 Nov 2021 16:32:06 +0530 Subject: [PATCH] Render the progress bar for file receives --- kittens/transfer/receive.py | 107 +++++++++++++++++++++++++++++++++--- 1 file changed, 99 insertions(+), 8 deletions(-) diff --git a/kittens/transfer/receive.py b/kittens/transfer/receive.py index 4289d5ecd..f4c0c2cfc 100644 --- a/kittens/transfer/receive.py +++ b/kittens/transfer/receive.py @@ -3,6 +3,7 @@ import os import posixpath +from asyncio import TimerHandle from collections import deque from contextlib import suppress from enum import auto @@ -11,7 +12,7 @@ from time import monotonic from typing import Deque, Dict, Iterator, List, Optional from kitty.cli_stub import TransferCLIOptions -from kitty.fast_data_types import FILE_TRANSFER_CODE +from kitty.fast_data_types import FILE_TRANSFER_CODE, wcswidth from kitty.file_transmission import ( Action, Compression, FileTransmissionCommand, FileType, NameReprEnum, encode_bypass @@ -22,9 +23,13 @@ from kitty.utils import sanitize_control_codes from ..tui.handler import Handler from ..tui.loop import Loop, debug from ..tui.operations import styled, without_line_wrap +from ..tui.spinners import Spinner from ..tui.utils import human_size from .send import Transfer -from .utils import expand_home, random_id, should_be_compressed +from .utils import ( + expand_home, random_id, render_progress_in_width, safe_divide, + should_be_compressed +) debug file_counter = count(1) @@ -172,6 +177,7 @@ class ProgressTracker: self.transfered_stats_interval = 0. self.started_at = 0. self.signature_bytes = 0 + self.done_files: List[File] = [] def change_active_file(self, nf: File) -> None: now = monotonic() @@ -183,7 +189,8 @@ class ProgressTracker: self.started_at = monotonic() def file_written(self, af: File, amt: int, is_done: bool) -> None: - self.active_file = af + if self.active_file is not af: + self.change_active_file(af) af.transmitted_bytes += amt self.total_transferred += amt self.transfers.append(Transfer(amt)) @@ -194,6 +201,7 @@ class ProgressTracker: self.transfered_stats_amt = sum(t.amt for t in self.transfers) if is_done: af.done_at = monotonic() + self.done_files.append(af) class Manager: @@ -334,6 +342,10 @@ class Receive(Handler): self.quit_after_write_code: Optional[int] = None self.check_paths_printed = False self.transmit_started = False + self.max_name_length = 0 + self.spinner = Spinner() + self.progress_update_call: Optional[TimerHandle] = None + self.progress_drawn = False def send_payload(self, payload: str) -> None: self.write(self.manager.prefix) @@ -382,10 +394,10 @@ class Receive(Handler): else: self.start_transfer() if self.manager.transfer_done: - self.exit_after_completion() - - def exit_after_completion(self) -> None: - self.quit_loop(0) + self.quit_after_write_code = 0 + self.refresh_progress() + elif self.transmit_started: + self.refresh_progress() def confirm_paths(self) -> None: self.print_check_paths() @@ -438,6 +450,8 @@ class Receive(Handler): self.print(f'Queueing transfer of {len(self.manager.files)} files(s)') for x in self.manager.request_files(): self.send_payload(x) + names = (f.display_name for f in self.manager.files) + self.max_name_length = max(6, max(map(wcswidth, names))) def print_err(self, msg: str) -> None: self.cmd.styled(msg, fg='red') @@ -463,12 +477,89 @@ class Receive(Handler): self.manager.state = State.canceled self.asyncio_loop.call_later(delay, self.quit_loop, 1) + def render_progress( + self, name: str, spinner_char: str = ' ', bytes_so_far: int = 0, total_bytes: int = 0, + secs_so_far: float = 0., bytes_per_sec: float = 0., is_complete: bool = False + ) -> None: + if is_complete: + bytes_so_far = total_bytes + self.write(render_progress_in_width( + name, width=self.screen_size.cols, max_path_length=self.max_name_length, spinner_char=spinner_char, + bytes_so_far=bytes_so_far, total_bytes=total_bytes, secs_so_far=secs_so_far, + bytes_per_sec=bytes_per_sec, is_complete=is_complete + )) + + def draw_progress_for_current_file(self, af: File, spinner_char: str = ' ', is_complete: bool = False) -> None: + p = self.manager.progress_tracker + now = monotonic() + self.render_progress( + af.display_name, spinner_char=spinner_char, is_complete=is_complete, + bytes_so_far=af.transmitted_bytes, total_bytes=af.expected_size, + secs_so_far=(af.done_at or now) - af.transmit_started_at, + bytes_per_sec=safe_divide(p.transfered_stats_amt, p.transfered_stats_interval) + ) + + def erase_progress(self) -> None: + if self.progress_drawn: + self.cmd.move_cursor_by(2, 'up') + self.write('\r') + self.cmd.clear_to_end_of_screen() + self.progress_drawn = False + + def refresh_progress(self) -> None: + self.erase_progress() + self.draw_progress() + + def schedule_progress_update(self, delay: float = 0.1) -> None: + if self.progress_update_call is None: + self.progress_update_call = self.asyncio_loop.call_later(delay, self.refresh_progress) + elif self.asyncio_loop.time() + delay < self.progress_update_call.when(): + self.progress_update_call.cancel() + self.progress_update_call = self.asyncio_loop.call_later(delay, self.refresh_progress) + @Handler.atomic_update def draw_progress(self) -> None: if self.manager.state is State.canceled: return with without_line_wrap(self.write): - pass + for df in self.manager.progress_tracker.done_files: + sc = styled('✔', fg='green') + if df.ftype is FileType.regular: + self.draw_progress_for_current_file(df, spinner_char=sc, is_complete=True) + else: + self.write(f'{sc} {df.display_name} {styled(df.ftype.name, dim=True, italic=True)}') + self.print() + del self.manager.progress_tracker.done_files[:] + is_complete = self.quit_after_write_code is not None + if is_complete: + sc = styled('✔', fg='green') if self.quit_after_write_code == 0 else styled('✘', fg='red') + else: + sc = self.spinner() + p = self.manager.progress_tracker + now = monotonic() + if is_complete: + self.cmd.repeat('─', self.screen_size.width) + else: + af = p.active_file + if af is not None: + self.draw_progress_for_current_file(af, spinner_char=sc) + self.print() + if p.total_transferred > 0: + self.render_progress( + 'Total', spinner_char=sc, + bytes_so_far=p.total_transferred, total_bytes=p.total_bytes_to_transfer, + secs_so_far=now - p.started_at, is_complete=is_complete, + bytes_per_sec=safe_divide(p.transfered_stats_amt, p.transfered_stats_interval) + ) + else: + self.print('File data transfer has not yet started', end='') + self.print() + self.schedule_progress_update(self.spinner.interval) + self.progress_drawn = True + + def on_writing_finished(self) -> None: + if self.quit_after_write_code is not None: + self.quit_loop(self.quit_after_write_code) def receive_main(cli_opts: TransferCLIOptions, args: List[str]) -> None: