From 820a893d75e641ca2a90e99bf3d7c2a4d70421d6 Mon Sep 17 00:00:00 2001 From: Kovid Goyal Date: Fri, 29 Oct 2021 15:01:31 +0530 Subject: [PATCH] More work on file transfer --- kittens/transfer/receive.py | 38 ++++++++++++++++++++++---------- kitty_tests/file_transmission.py | 38 +++++++++++++++++++++++++++++--- 2 files changed, 61 insertions(+), 15 deletions(-) diff --git a/kittens/transfer/receive.py b/kittens/transfer/receive.py index b7a3432c0..d20998bcc 100644 --- a/kittens/transfer/receive.py +++ b/kittens/transfer/receive.py @@ -34,31 +34,42 @@ class File: self.mtime = ftc.mtime self.spec_id = int(ftc.file_id) self.permissions = ftc.permissions - self.remote_name = ftc.name + self.remote_path = ftc.name self.remote_id = ftc.status self.remote_target = ftc.data.decode('utf-8') self.parent = ftc.parent - self.local_name = '' + self.expanded_local_path = '' + + def __repr__(self) -> str: + return f'File(rpath={self.remote_path!r}, lpath={self.expanded_local_path!r})' class TreeNode: def __init__(self, file: File, local_name: str, parent: Optional['TreeNode'] = None): self.entry = file - self.entry.local_name = local_name - self.children: List[TreeNode] = [] + self.entry.expanded_local_path = local_name self.parent = parent + self.added_files: Dict[int, TreeNode] = {} def add_child(self, file: File) -> 'TreeNode': - c = TreeNode(file, os.path.join(self.entry.local_name, os.path.basename(file.remote_name)), self) - self.children.append(c) + q = self.added_files.get(id(file)) + if q is not None: + return q + c = TreeNode(file, os.path.join(self.entry.expanded_local_path, os.path.basename(file.remote_path)), self) + self.added_files[id(file)] = c return c + def __iter__(self) -> Iterator['TreeNode']: + for c in self.added_files.values(): + yield c + yield from c + def make_tree(all_files: List[File], local_base: str) -> TreeNode: fid_map = {f.remote_id: f for f in all_files} node_map: Dict[str, TreeNode] = {} - root_node = TreeNode(File(FileTransmissionCommand()), local_base) + root_node = TreeNode(File(FileTransmissionCommand(file_id='-1')), local_base) def ensure_parent(f: File) -> TreeNode: if not f.parent: @@ -76,11 +87,11 @@ def make_tree(all_files: List[File], local_base: str) -> TreeNode: return root_node -def files_for_receive(cli_opts: TransferCLIOptions, dest: str, files: List[File], remote_home: str, specs: List[str]) -> None: +def files_for_receive(cli_opts: TransferCLIOptions, dest: str, files: List[File], remote_home: str, specs: List[str]) -> Iterator[File]: spec_map: Dict[int, List[File]] = {i: [] for i in range(len(specs))} for f in files: spec_map[f.spec_id].append(f) - spec_paths = [spec_map[i][0].remote_name for i in range(len(specs))] + spec_paths = [spec_map[i][0].remote_path for i in range(len(specs))] if cli_opts.mode == 'mirror': try: common_path = posixpath.commonpath(spec_paths) @@ -91,16 +102,19 @@ def files_for_receive(cli_opts: TransferCLIOptions, dest: str, files: List[File] spec_paths = [posixpath.join('~', posixpath.relpath(x, home)) for x in spec_paths] for spec_id, files_for_spec in spec_map.items(): spec = spec_paths[spec_id] - tree = make_tree(files_for_spec, os.path.expanduser(spec)) + tree = make_tree(files_for_spec, os.path.dirname(os.path.expanduser(spec))) + for x in tree: + yield x.entry else: dest_is_dir = dest[-1] in (os.sep, os.altsep) or len(specs) > 1 for spec_id, files_for_spec in spec_map.items(): if dest_is_dir: - dest_path = os.path.join(dest, posixpath.basename(files_for_spec[0].remote_name)) + dest_path = os.path.join(dest, posixpath.basename(files_for_spec[0].remote_path)) else: dest_path = dest tree = make_tree(files_for_spec, os.path.expanduser(dest_path)) - tree + for x in tree: + yield x.entry class Manager: diff --git a/kitty_tests/file_transmission.py b/kitty_tests/file_transmission.py index 45c0747bd..e0f816281 100644 --- a/kitty_tests/file_transmission.py +++ b/kitty_tests/file_transmission.py @@ -13,12 +13,14 @@ from kittens.transfer.librsync import ( LoadSignature, PatchFile, delta_for_file, signature_of_file ) from kittens.transfer.main import parse_transfer_args +from kittens.transfer.receive import files_for_receive, File from kittens.transfer.rsync import decode_utf8_buffer, parse_ftc from kittens.transfer.send import files_for_send -from kittens.transfer.utils import set_paths +from kittens.transfer.utils import set_paths, home_path from kitty.file_transmission import ( Action, Compression, FileTransmissionCommand, FileType, - TestFileTransmission as FileTransmission, TransmissionType + TestFileTransmission as FileTransmission, TransmissionType, + iter_file_metadata ) from . import BaseTest @@ -345,7 +347,37 @@ class TestFileTransmission(BaseTest): t('a1=b1;c=d;;e', 'a1', 'b1', 'c', 'd;e') t('a1=b1;c=d;;;1=1', 'a1', 'b1', 'c', 'd;', '1', '1') - def test_path_mapping(self): + def test_path_mapping_receive(self): + opts = parse_transfer_args([])[0] + b = Path(os.path.join(self.tdir, 'b')) + os.makedirs(b) + open(b / 'r', 'w').close() + os.mkdir(b / 'd') + open(b / 'd' / 'r', 'w').close() + + def am(files, kw): + m = {f.remote_path: f.expanded_local_path for f in files} + kw = {str(k): str(v) for k, v in kw.items()} + self.ae(kw, m) + + def tf(args, expected): + if opts.mode == 'mirror': + all_specs = args + dest = '' + else: + all_specs = args[:-1] + dest = args[-1] + specs = list((str(i), str(s)) for i, s in enumerate(all_specs)) + files = list(map(File, iter_file_metadata(specs))) + files = list(files_for_receive(opts, dest, files, home_path(), specs)) + self.ae(len(files), len(expected)) + am(files, expected) + + opts.mode = 'mirror' + with set_paths(cwd=b, home='/foo/bar'): + tf([b/'r', b/'d'], {b/'r': b/'r', b/'d': b/'d', b/'d'/'r': b/'d'/'r'}) + + def test_path_mapping_send(self): opts = parse_transfer_args([])[0] b = Path(os.path.join(self.tdir, 'b')) os.makedirs(b)