diff --git a/kittens/transfer/receive.py b/kittens/transfer/receive.py index 07f63e972..4289d5ecd 100644 --- a/kittens/transfer/receive.py +++ b/kittens/transfer/receive.py @@ -3,10 +3,12 @@ import os import posixpath +from collections import deque from contextlib import suppress from enum import auto from itertools import count -from typing import Dict, Iterator, List, Optional +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 @@ -21,6 +23,7 @@ from ..tui.handler import Handler from ..tui.loop import Loop, debug from ..tui.operations import styled, without_line_wrap from ..tui.utils import human_size +from .send import Transfer from .utils import expand_home, random_id, should_be_compressed debug @@ -38,6 +41,8 @@ class File: def __init__(self, ftc: FileTransmissionCommand): self.expected_size = ftc.size + self.transmit_started_at = self.done_at = 0. + self.transmitted_bytes = 0 self.ftype = ftc.ftype self.mtime = ftc.mtime self.spec_id = int(ftc.file_id) @@ -56,10 +61,10 @@ class File: def __repr__(self) -> str: return f'File(rpath={self.remote_path!r}, lpath={self.expanded_local_path!r})' - def write_data(self, data: bytes, is_last: bool) -> None: + def write_data(self, data: bytes, is_last: bool) -> int: if self.ftype is FileType.symlink: self.remote_symlink_value += data - return + return 0 if self.ftype is FileType.regular: if not self.local_write_started: parent = os.path.dirname(self.expanded_local_path) @@ -67,7 +72,10 @@ class File: os.makedirs(parent, exist_ok=True) self.local_write_started = True with open(self.expanded_local_path, 'ab') as f: + base = f.tell() f.write(data) + return f.tell() - base + return 0 def apply_metadata(self) -> None: if self.ftype is FileType.symlink: @@ -152,6 +160,42 @@ def files_for_receive(cli_opts: TransferCLIOptions, dest: str, files: List[File] yield x.entry +class ProgressTracker: + + def __init__(self) -> None: + self.total_size_of_all_files = 0 + self.total_bytes_to_transfer = 0 + self.active_file: Optional[File] = None + self.total_transferred = 0 + self.transfers: Deque[Transfer] = deque() + self.transfered_stats_amt = 0 + self.transfered_stats_interval = 0. + self.started_at = 0. + self.signature_bytes = 0 + + def change_active_file(self, nf: File) -> None: + now = monotonic() + self.active_file = nf + nf.transmit_started_at = now + + def start_transfer(self) -> None: + self.transfers.append(Transfer()) + self.started_at = monotonic() + + def file_written(self, af: File, amt: int, is_done: bool) -> None: + self.active_file = af + af.transmitted_bytes += amt + self.total_transferred += amt + self.transfers.append(Transfer(amt)) + now = self.transfers[-1].at + while len(self.transfers) > 2 and self.transfers[0].is_too_old(now): + self.transfers.popleft() + self.transfered_stats_interval = now - self.transfers[0].at + self.transfered_stats_amt = sum(t.amt for t in self.transfers) + if is_done: + af.done_at = monotonic() + + class Manager: def __init__( @@ -169,13 +213,14 @@ class Manager: self.suffix = '\x1b\\' self.state = State.waiting_for_permission self.files: List[File] = [] - self.total_transfer_size = 0 + self.progress_tracker = ProgressTracker() self.transfer_done = False def start_transfer(self) -> Iterator[str]: yield FileTransmissionCommand(action=Action.receive, bypass=self.bypass, size=len(self.spec)).serialize() for i, x in enumerate(self.spec): yield FileTransmissionCommand(action=Action.file, file_id=str(i), name=x).serialize() + self.progress_tracker.start_transfer() def finalize_transfer(self) -> str: self.transfer_done = True @@ -221,7 +266,8 @@ class Manager: def collect_files(self, cli_opts: TransferCLIOptions) -> None: self.files = list(files_for_receive(cli_opts, self.dest, self.files, self.remote_home, self.spec)) - self.total_transfer_size = sum(max(0, f.expected_size) for f in self.files) + self.progress_tracker.total_size_of_all_files = sum(max(0, f.expected_size) for f in self.files) + self.progress_tracker.total_bytes_to_transfer = self.progress_tracker.total_size_of_all_files self.fid_map = {f.file_id: f for f in self.files} def on_file_transfer_response(self, ftc: FileTransmissionCommand) -> str: @@ -266,11 +312,13 @@ class Manager: f = self.fid_map.get(ftc.file_id) if f is None: return f'Got data for unknown file id: {ftc.file_id}' + is_last = ftc.action is Action.end_data try: - f.write_data(ftc.data, ftc.action is Action.end_data) + amt_written = f.write_data(ftc.data, is_last) except OSError as err: return str(err) - if ftc.action is Action.end_data: + self.progress_tracker.file_written(f, amt_written, is_last) + if is_last: del self.fid_map[ftc.file_id] if not self.fid_map: return self.finalize_transfer() @@ -353,7 +401,7 @@ class Receive(Handler): self.print(df.display_name, '→', end=' ') self.cmd.styled(df.expanded_local_path, fg='red' if os.path.lexists(df.expanded_local_path) else None) self.print() - self.print(f'Transferring {len(self.manager.files)} file(s) of total size: {human_size(self.manager.total_transfer_size)}') + self.print(f'Transferring {len(self.manager.files)} file(s) of total size: {human_size(self.manager.progress_tracker.total_size_of_all_files)}') self.print() self.print_continue_msg() @@ -417,6 +465,8 @@ class Receive(Handler): @Handler.atomic_update def draw_progress(self) -> None: + if self.manager.state is State.canceled: + return with without_line_wrap(self.write): pass