From 5cb36a463296b79bd1946e172838e9ef975668c2 Mon Sep 17 00:00:00 2001
From: Kovid Goyal <kovid@kovidgoyal.net>
Date: Tue, 26 Oct 2021 22:05:38 +0530
Subject: [PATCH] More work on file transfer

---
 kittens/transfer/receive.py | 66 +++++++++++++++++++++++++++++++++++++
 kitty/file_transmission.py  |  7 ++--
 2 files changed, 70 insertions(+), 3 deletions(-)

diff --git a/kittens/transfer/receive.py b/kittens/transfer/receive.py
index e062d03be..10fc49dbc 100644
--- a/kittens/transfer/receive.py
+++ b/kittens/transfer/receive.py
@@ -1,6 +1,8 @@
 #!/usr/bin/env python
 # License: GPLv3 Copyright: 2021, Kovid Goyal <kovid at kovidgoyal.net>
 
+import os
+import posixpath
 from enum import auto
 from typing import Dict, Iterator, List, Optional
 
@@ -35,6 +37,70 @@ class File:
         self.remote_name = ftc.name
         self.remote_id = ftc.status
         self.remote_target = ftc.data.decode('utf-8')
+        self.parent = ftc.parent
+        self.local_name = ''
+
+
+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.parent = parent
+
+    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)
+        return 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)
+
+    def ensure_parent(f: File) -> TreeNode:
+        if not f.parent:
+            return root_node
+        parent = node_map.get(f.parent)
+        if parent is None:
+            fp = fid_map[f.parent]
+            gp = ensure_parent(fp)
+            parent = gp.add_child(fp)
+        return parent
+
+    for f in all_files:
+        p = ensure_parent(f)
+        p.add_child(f)
+    return root_node
+
+
+def files_for_receive(cli_opts: TransferCLIOptions, dest: str, files: List[File], remote_home: str, specs: List[str]) -> None:
+    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))]
+    if cli_opts.mode == 'mirror':
+        try:
+            common_path = posixpath.commonpath(spec_paths)
+        except ValueError:
+            common_path = ''
+        home = remote_home.rstrip('/')
+        if common_path and common_path.startswith(home + '/'):
+            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))
+    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))
+            else:
+                dest_path = dest
+            tree = make_tree(files_for_spec, os.path.expanduser(dest_path))
+    tree
 
 
 class Manager:
diff --git a/kitty/file_transmission.py b/kitty/file_transmission.py
index f29f861dc..cc574dc80 100644
--- a/kitty/file_transmission.py
+++ b/kitty/file_transmission.py
@@ -67,7 +67,7 @@ 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)
 
-    def make_ftc(path: str, spec_id: str, sr: Optional[os.stat_result] = None) -> FileTransmissionCommand:
+    def make_ftc(path: str, spec_id: str, sr: Optional[os.stat_result] = None, parent: str = '') -> FileTransmissionCommand:
         if sr is None:
             sr = os.stat(path, follow_symlinks=False)
         if stat.S_ISLNK(sr.st_mode):
@@ -80,7 +80,7 @@ def iter_file_metadata(file_specs: Iterable[Tuple[str, str]]) -> Iterator[Union[
             raise ValueError('Not an appropriate file type')
         return 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
+            name=path, status=f'{sr.st_dev}:{sr.st_ino}', size=sr.st_size, ftype=ftype, parent=parent
         )
 
     for spec_id, spec in file_specs:
@@ -115,7 +115,7 @@ 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)
+                            dftc = make_ftc(os.path.join(dirpath, dname), spec_id, parent=ftc.status)
                         except (ValueError, OSError):
                             continue
                         file_map[dftc.status].append(dftc)
@@ -235,6 +235,7 @@ class FileTransmissionCommand:
     size: int = -1
     name: str = field(default='', metadata={'base64': True})
     status: str = field(default='', metadata={'base64': True})
+    parent: str = field(default='', metadata={'base64': True})
     data: bytes = field(default=b'', repr=False)
 
     def __repr__(self) -> str: