More work on file transmission

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]
+                    del self.files_to_be_transferred[ftc.file_id]
-                    if not self.fid_map:
+                    if not self.files_to_be_transferred:
                         return self.finalize_transfer()
         return ''
 

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():
+    def resolve_symlink(ftc: FileTransmissionCommand) -> FileTransmissionCommand:
-        base = cmds[0]
+        if ftc.ftype is FileType.symlink:
-        if base.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}'
+                    tgt = file_map.get(skey(s))
-                    if fkey in file_map:
+                    if tgt is not None:
-                        base.data = fkey.encode('utf-8', 'replace')
+                        ftc.data = tgt[0].status.encode('utf-8')
-        yield base
+        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:]:
+            for q in cmds:
-                if q.ftype is FileType.regular:
+                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