From a1ca607f35c1dcea7f47814a52351c0d6b0de9ec Mon Sep 17 00:00:00 2001 From: Kovid Goyal Date: Thu, 11 Nov 2021 20:05:54 +0530 Subject: [PATCH] More work on file transmission --- kittens/transfer/receive.py | 23 +++++++++++++------- kitty/file_transmission.py | 42 ++++++++++++++++++++++--------------- 2 files changed, 40 insertions(+), 25 deletions(-) diff --git a/kittens/transfer/receive.py b/kittens/transfer/receive.py index f4c0c2cfc..707128ba5 100644 --- a/kittens/transfer/receive.py +++ b/kittens/transfer/receive.py @@ -238,16 +238,18 @@ class Manager: try: os.makedirs(f.expanded_local_path, exist_ok=True) except OSError as err: - return str(err) + return f'Failed to create directory with error: {err}' elif f.ftype is FileType.link: target = rid_map.get(f.remote_target) if target is None: return f'Hard link with remote id: {f.remote_target} not found' try: os.makedirs(os.path.dirname(f.expanded_local_path), exist_ok=True) + with suppress(FileNotFoundError): + os.remove(f.expanded_local_path) os.link(target.expanded_local_path, f.expanded_local_path) except OSError as err: - return str(err) + return f'Failed to create hardlink with error: {err}' elif f.ftype is FileType.symlink: if f.remote_target: target = rid_map.get(f.remote_target) @@ -258,7 +260,12 @@ class Manager: lt = os.path.relpath(lt, os.path.dirname(f.expanded_local_path)) else: lt = f.remote_symlink_value.decode('utf-8') - os.symlink(lt, f.expanded_local_path) + with suppress(FileNotFoundError): + os.remove(f.expanded_local_path) + try: + os.symlink(lt, f.expanded_local_path) + except OSError as err: + return f'Failed to create symlink with error: {err}' with suppress(OSError): f.apply_metadata() return '' @@ -274,9 +281,9 @@ 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.progress_tracker.total_size_of_all_files = sum(max(0, f.expected_size) for f in self.files) + self.files_to_be_transferred = {f.file_id: f for f in self.files if f.ftype not in (FileType.directory, FileType.link)} + self.progress_tracker.total_size_of_all_files = sum(max(0, f.expected_size) for f in self.files_to_be_transferred.values()) 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: if self.state is State.waiting_for_permission: @@ -317,7 +324,7 @@ class Manager: return f'Unexpected response from terminal: {ftc}' elif self.state is State.transferring: if ftc.action in (Action.data, Action.end_data): - f = self.fid_map.get(ftc.file_id) + f = self.files_to_be_transferred.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 @@ -327,8 +334,8 @@ class Manager: return str(err) 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: + del self.files_to_be_transferred[ftc.file_id] + if not self.files_to_be_transferred: return self.finalize_transfer() return '' diff --git a/kitty/file_transmission.py b/kitty/file_transmission.py index bdcfe1b14..bb64ad98d 100644 --- a/kitty/file_transmission.py +++ b/kitty/file_transmission.py @@ -12,7 +12,7 @@ from dataclasses import Field, dataclass, field, fields from enum import Enum, auto from functools import partial from gettext import gettext as _ -from itertools import chain +from itertools import chain, count from time import monotonic from typing import ( IO, Any, Callable, DefaultDict, Deque, Dict, Iterable, Iterator, List, @@ -67,7 +67,11 @@ def split_for_transfer( def iter_file_metadata(file_specs: Iterable[Tuple[str, str]]) -> Iterator[Union['FileTransmissionCommand', 'TransmissionError']]: - file_map: DefaultDict[str, List[FileTransmissionCommand]] = defaultdict(list) + file_map: DefaultDict[Tuple[int, int], List[FileTransmissionCommand]] = defaultdict(list) + counter = count() + + def skey(sr: os.stat_result) -> Tuple[int, int]: + return sr.st_dev, sr.st_ino def make_ftc(path: str, spec_id: str, sr: Optional[os.stat_result] = None, parent: str = '') -> FileTransmissionCommand: if sr is None: @@ -80,10 +84,12 @@ def iter_file_metadata(file_specs: Iterable[Tuple[str, str]]) -> Iterator[Union[ ftype = FileType.regular else: raise ValueError('Not an appropriate file type') - return FileTransmissionCommand( + ans = FileTransmissionCommand( action=Action.file, file_id=spec_id, mtime=sr.st_mtime_ns, permissions=stat.S_IMODE(sr.st_mode), - name=path, status=f'{sr.st_dev}:{sr.st_ino}', size=sr.st_size, ftype=ftype, parent=parent + name=path, status=str(next(counter)), size=sr.st_size, ftype=ftype, parent=parent ) + file_map[skey(sr)].append(ans) + return ans for spec_id, spec in file_specs: path = spec @@ -106,7 +112,6 @@ def iter_file_metadata(file_specs: Iterable[Tuple[str, str]]) -> Iterator[Union[ except ValueError: yield TransmissionError(file_id=spec_id, code='EINVAL', msg='Not a valid filetype') continue - file_map[ftc.status].append(ftc) if ftc.ftype is FileType.directory: try: x_ok = os.access(path, os.X_OK) @@ -117,15 +122,14 @@ def iter_file_metadata(file_specs: Iterable[Tuple[str, str]]) -> Iterator[Union[ for dirpath, dirnames, filenames in di: for dname in chain(dirnames, filenames): try: - dftc = make_ftc(os.path.join(dirpath, dname), spec_id, parent=ftc.status) + make_ftc(os.path.join(dirpath, dname), spec_id, parent=ftc.status) except (ValueError, OSError): continue - file_map[dftc.status].append(dftc) - for fkey, cmds in file_map.items(): - base = cmds[0] - if base.ftype is FileType.symlink: + + def resolve_symlink(ftc: FileTransmissionCommand) -> FileTransmissionCommand: + if ftc.ftype is FileType.symlink: try: - dest = os.path.realpath(base.name) + dest = os.path.realpath(ftc.name) except OSError: pass else: @@ -134,13 +138,17 @@ def iter_file_metadata(file_specs: Iterable[Tuple[str, str]]) -> Iterator[Union[ except OSError: pass else: - fkey = f'{s.st_dev}:{s.st_ino}' - if fkey in file_map: - base.data = fkey.encode('utf-8', 'replace') - yield base + tgt = file_map.get(skey(s)) + if tgt is not None: + ftc.data = tgt[0].status.encode('utf-8') + return ftc + + for fkey, cmds in file_map.items(): + base = cmds[0] + yield resolve_symlink(base) if len(cmds) > 1 and base.ftype is FileType.regular: - for q in cmds[1:]: - if q.ftype is FileType.regular: + for q in cmds: + if q is not base and q.ftype is FileType.regular: q.ftype = FileType.link q.data = base.status.encode('utf-8', 'replace') yield q