More work on transmission kitten
This commit is contained in:
Kovid Goyal
parent
d7a1f3104a
commit
af5c0b3b71
@ -2,11 +2,47 @@
|
||||
# vim:fileencoding=utf-8
|
||||
# License: GPLv3 Copyright: 2021, Kovid Goyal <kovid at kovidgoyal.net>
|
||||
|
||||
import os
|
||||
import stat
|
||||
import sys
|
||||
from typing import List
|
||||
from contextlib import contextmanager
|
||||
from itertools import count
|
||||
from typing import (
|
||||
Dict, Generator, Iterable, Iterator, List, Sequence, Tuple, Union
|
||||
)
|
||||
|
||||
from kitty.cli import parse_args
|
||||
from kitty.cli_stub import TransferCLIOptions
|
||||
from kitty.file_transmission import (
|
||||
Action, Compression, FileTransmissionCommand, FileType
|
||||
)
|
||||
|
||||
_cwd = _home = ''
|
||||
|
||||
|
||||
def abspath(path: str) -> str:
|
||||
return os.path.normpath(os.path.join(_cwd or os.getcwd(), path))
|
||||
|
||||
|
||||
def home_path() -> str:
|
||||
return _home or os.path.expanduser('~')
|
||||
|
||||
|
||||
def expand_home(path: str) -> str:
|
||||
if path.startswith('~' + os.sep) or (os.altsep and path.startswith('~' + os.altsep)):
|
||||
return os.path.join(home_path(), path[2:].lstrip(os.sep + (os.altsep or '')))
|
||||
return path
|
||||
|
||||
|
||||
@contextmanager
|
||||
def set_paths(cwd: str = '', home: str = '') -> Generator[None, None, None]:
|
||||
global _cwd, _home
|
||||
orig = _cwd, _home
|
||||
try:
|
||||
_cwd, _home = cwd, home
|
||||
yield
|
||||
finally:
|
||||
_cwd, _home = orig
|
||||
|
||||
|
||||
def option_text() -> str:
|
||||
@ -15,18 +51,190 @@ def option_text() -> str:
|
||||
default=send
|
||||
choices=send,receive
|
||||
Whether to send or receive files.
|
||||
|
||||
|
||||
--mode -m
|
||||
default=normal
|
||||
choices=mirror
|
||||
How to interpret command line arguments. In :code:`mirror` mode all arguments
|
||||
are assumed to be files on the sending computer and they are mirrored onto the
|
||||
receiving computer. In :code:`normal` mode the last argument is assumed to be a
|
||||
destination path on the receiving computer.
|
||||
'''
|
||||
|
||||
|
||||
class IdentityCompressor:
|
||||
|
||||
def compress(self, data: bytes) -> bytes:
|
||||
return data
|
||||
|
||||
def flush(self) -> bytes:
|
||||
return b''
|
||||
|
||||
|
||||
class ZlibCompressor:
|
||||
|
||||
def __init__(self) -> None:
|
||||
import zlib
|
||||
self.c = zlib.compressobj()
|
||||
|
||||
def compress(self, data: bytes) -> bytes:
|
||||
return self.c.compress(data)
|
||||
|
||||
def flush(self) -> bytes:
|
||||
return self.c.flush()
|
||||
|
||||
|
||||
def get_remote_path(local_path: str, remote_base: str) -> str:
|
||||
if not remote_base:
|
||||
return local_path.replace(os.sep, '/')
|
||||
if remote_base.endswith('/'):
|
||||
return os.path.join(remote_base, os.path.basename(local_path))
|
||||
return remote_base
|
||||
|
||||
|
||||
class File:
|
||||
|
||||
def __init__(
|
||||
self, local_path: str, expanded_local_path: str, file_id: int, stat_result: os.stat_result, remote_base: str, file_type: FileType
|
||||
) -> None:
|
||||
self.local_path = local_path
|
||||
self.expanded_local_path = expanded_local_path
|
||||
self.permissions = stat.S_IMODE(stat_result.st_mode)
|
||||
self.mtime = stat_result.st_mtime_ns
|
||||
self.file_size = stat_result.st_size
|
||||
self.file_hash = stat_result.st_dev, stat_result.st_ino
|
||||
self.remote_path = get_remote_path(self.local_path, remote_base)
|
||||
self.remote_path = self.remote_path.replace(os.sep, '/')
|
||||
self.file_id = hex(file_id)[2:]
|
||||
self.hard_link_target = ''
|
||||
self.symbolic_link_target = ''
|
||||
self.stat_result = stat_result
|
||||
self.file_type = file_type
|
||||
self.compression = Compression.zlib if self.file_size > 2048 else Compression.none
|
||||
|
||||
def metadata_command(self) -> FileTransmissionCommand:
|
||||
return FileTransmissionCommand(
|
||||
action=Action.file, compression=self.compression, ftype=self.file_type,
|
||||
name=self.remote_path, permissions=self.permissions, mtime=self.mtime,
|
||||
file_id=self.file_id,
|
||||
)
|
||||
|
||||
def data_commands(self) -> Iterator[FileTransmissionCommand]:
|
||||
if self.file_type is FileType.symlink:
|
||||
yield FileTransmissionCommand(action=Action.end_data, data=self.symbolic_link_target.encode('utf-8'), file_id=self.file_id)
|
||||
elif self.file_type is FileType.link:
|
||||
yield FileTransmissionCommand(action=Action.end_data, data=f'fid:{self.hard_link_target}'.encode('utf-8'), file_id=self.file_id)
|
||||
elif self.file_type is FileType.regular:
|
||||
compressor: Union[IdentityCompressor, ZlibCompressor] = ZlibCompressor() if self.compression is Compression.zlib else IdentityCompressor()
|
||||
with open(self.local_path, 'rb') as f:
|
||||
keep_going = True
|
||||
while keep_going:
|
||||
data = f.read(4096)
|
||||
keep_going = bool(data)
|
||||
data = compressor.compress(data) if data else compressor.flush()
|
||||
if data or not keep_going:
|
||||
yield FileTransmissionCommand(action=Action.data if keep_going else Action.end_data, data=data, file_id=self.file_id)
|
||||
|
||||
|
||||
def process(cli_opts: TransferCLIOptions, paths: Iterable[str], remote_base: str) -> Iterator[File]:
|
||||
counter = count(1)
|
||||
for x in paths:
|
||||
expanded = expand_home(x)
|
||||
try:
|
||||
s = os.stat(expanded, follow_symlinks=False)
|
||||
except OSError as err:
|
||||
raise SystemExit(f'Failed to stat {x} with error: {err}') from err
|
||||
if stat.S_ISDIR(s.st_mode):
|
||||
yield File(x, expanded, next(counter), s, remote_base, FileType.directory)
|
||||
new_remote_base = remote_base
|
||||
if new_remote_base:
|
||||
new_remote_base = new_remote_base.rstrip('/') + '/' + os.path.basename(x) + '/'
|
||||
else:
|
||||
new_remote_base = x.replace(os.sep, '/').rstrip('/') + '/'
|
||||
yield from process(cli_opts, [os.path.join(x, y) for y in os.listdir(expanded)], new_remote_base)
|
||||
elif stat.S_ISLNK(s.st_mode):
|
||||
yield File(x, expanded, next(counter), s, remote_base, FileType.symlink)
|
||||
elif stat.S_ISREG(s.st_mode):
|
||||
yield File(x, expanded, next(counter), s, remote_base, FileType.regular)
|
||||
|
||||
|
||||
def process_mirrored_files(cli_opts: TransferCLIOptions, args: Sequence[str]) -> Iterator[File]:
|
||||
paths = [abspath(x) for x in args]
|
||||
try:
|
||||
common_path = os.path.commonpath(paths)
|
||||
except ValueError:
|
||||
common_path = ''
|
||||
home = home_path().rstrip(os.sep)
|
||||
if common_path and common_path.startswith(home + os.sep):
|
||||
paths = [os.path.join('~', os.path.relpath(x, home)) for x in paths]
|
||||
yield from process(cli_opts, paths, '')
|
||||
|
||||
|
||||
def process_normal_files(cli_opts: TransferCLIOptions, args: Sequence[str]) -> Iterator[File]:
|
||||
if len(args) < 2:
|
||||
raise SystemExit('Must specify at least one local path and one remote path')
|
||||
args = list(args)
|
||||
remote_base = args.pop().replace(os.sep, '/')
|
||||
if len(args) > 1 and not remote_base.endswith('/'):
|
||||
remote_base += '/'
|
||||
paths = [abspath(x) for x in args]
|
||||
yield from process(cli_opts, paths, remote_base)
|
||||
|
||||
|
||||
def files_for_send(cli_opts: TransferCLIOptions, args: List[str]) -> Tuple[File, ...]:
|
||||
if cli_opts.mode == 'mirror':
|
||||
files = list(process_mirrored_files(cli_opts, args))
|
||||
else:
|
||||
files = list(process_normal_files(cli_opts, args))
|
||||
groups: Dict[Tuple[int, int], List[File]] = {}
|
||||
|
||||
# detect hard links
|
||||
for f in files:
|
||||
groups.setdefault(f.file_hash, []).append(f)
|
||||
for group in groups.values():
|
||||
if len(group) > 1:
|
||||
for lf in group[1:]:
|
||||
lf.file_type = FileType.link
|
||||
lf.hard_link_target = group[0].file_id
|
||||
|
||||
# detect symlinks to other transferred files
|
||||
for f in tuple(files):
|
||||
if f.file_type is FileType.symlink:
|
||||
try:
|
||||
link_dest = os.readlink(f.local_path)
|
||||
except OSError:
|
||||
files.remove(f)
|
||||
continue
|
||||
f.symbolic_link_target = f'path:{link_dest}'
|
||||
q = link_dest if os.path.isabs(link_dest) else os.path.join(os.path.dirname(f.local_path), link_dest)
|
||||
try:
|
||||
st = os.stat(q)
|
||||
except OSError:
|
||||
pass
|
||||
else:
|
||||
fh = st.st_dev, st.st_ino
|
||||
if fh in groups:
|
||||
g = tuple(x for x in groups[fh] if os.path.samestat(st, x.stat_result))
|
||||
if g:
|
||||
t = g[0]
|
||||
f.symbolic_link_target = f'fid:{t.file_id}'
|
||||
return tuple(files)
|
||||
|
||||
|
||||
def send_main(cli_opts: TransferCLIOptions, args: List[str]) -> None:
|
||||
pass
|
||||
|
||||
|
||||
def main(args: List[str]) -> None:
|
||||
cli_opts, items = parse_args(
|
||||
def parse_transfer_args(args: List[str]) -> Tuple[TransferCLIOptions, List[str]]:
|
||||
return parse_args(
|
||||
args[1:], option_text, '', 'Transfer files over the TTY device',
|
||||
'kitty transfer', result_class=TransferCLIOptions
|
||||
)
|
||||
|
||||
|
||||
def main(args: List[str]) -> None:
|
||||
cli_opts, items = parse_transfer_args(args)
|
||||
if not items:
|
||||
raise SystemExit('Usage: kitty +kitten transfer file_or_directory ...')
|
||||
if cli_opts.direction == 'send':
|
||||
|
||||
@ -55,7 +55,6 @@ class FileType(NameReprEnum):
|
||||
|
||||
class TransmissionType(NameReprEnum):
|
||||
simple = auto()
|
||||
resume = auto()
|
||||
rsync = auto()
|
||||
|
||||
|
||||
@ -197,7 +196,9 @@ class DestFile:
|
||||
def __init__(self, ftc: FileTransmissionCommand) -> None:
|
||||
self.name = ftc.name
|
||||
if not os.path.isabs(self.name):
|
||||
self.name = os.path.join(tempfile.gettempdir(), self.name)
|
||||
self.name = os.path.expanduser(self.name)
|
||||
if not os.path.isabs(self.name):
|
||||
self.name = os.path.join(tempfile.gettempdir(), self.name)
|
||||
self.mtime = ftc.mtime
|
||||
self.file_id = ftc.file_id
|
||||
self.permissions = ftc.permissions
|
||||
|
||||
@ -8,7 +8,11 @@ import shutil
|
||||
import stat
|
||||
import tempfile
|
||||
import zlib
|
||||
from pathlib import Path
|
||||
|
||||
from kittens.transfer.main import (
|
||||
files_for_send, parse_transfer_args, set_paths
|
||||
)
|
||||
from kitty.file_transmission import (
|
||||
Action, Compression, FileTransmissionCommand, FileType,
|
||||
TestFileTransmission as FileTransmission, TransmissionType
|
||||
@ -51,14 +55,20 @@ class TestFileTransmission(BaseTest):
|
||||
def setUp(self):
|
||||
self.tdir = os.path.realpath(tempfile.mkdtemp())
|
||||
self.responses = []
|
||||
self.orig_home = os.environ.get('HOME')
|
||||
|
||||
def tearDown(self):
|
||||
shutil.rmtree(self.tdir)
|
||||
self.responses = []
|
||||
if self.orig_home is None:
|
||||
os.environ.pop('HOME')
|
||||
else:
|
||||
os.environ['HOME'] = self.orig_home
|
||||
|
||||
def clean_tdir(self):
|
||||
shutil.rmtree(self.tdir)
|
||||
self.tdir = os.path.realpath(tempfile.mkdtemp())
|
||||
self.responses = []
|
||||
|
||||
def assertResponses(self, ft, **kw):
|
||||
self.responses.append(response(**kw))
|
||||
@ -128,7 +138,6 @@ class TestFileTransmission(BaseTest):
|
||||
|
||||
# multi file send
|
||||
self.clean_tdir()
|
||||
self.responses = []
|
||||
ft = FileTransmission()
|
||||
dest = os.path.join(self.tdir, '2.bin')
|
||||
ft.handle_serialized_command(serialized_cmd(action='send'))
|
||||
@ -176,3 +185,51 @@ class TestFileTransmission(BaseTest):
|
||||
self.ae(os.stat(dest + 'd1').st_mtime_ns, 29)
|
||||
self.ae(os.stat(dest + 'd2').st_mtime_ns, 29)
|
||||
self.assertFalse(ft.active_receives)
|
||||
|
||||
def test_path_mapping(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 gm(*args):
|
||||
return files_for_send(opts, list(map(str, args)))
|
||||
|
||||
def am(files, kw):
|
||||
m = {f.expanded_local_path: f.remote_path for f in files}
|
||||
kw = {str(k): str(v) for k, v in kw.items()}
|
||||
self.ae(m, kw)
|
||||
|
||||
def tf(args, expected):
|
||||
files = gm(*args)
|
||||
self.ae(len(files), len(expected))
|
||||
am(files, expected)
|
||||
|
||||
opts.mode = 'mirror'
|
||||
with set_paths(cwd=b, home='/foo/bar'):
|
||||
tf(['r', 'd'], {b/'r': b/'r', b/'d': b/'d', b/'d'/'r': b/'d'/'r'})
|
||||
tf(['r', 'd/r'], {b/'r': b/'r', b/'d'/'r': b/'d'/'r'})
|
||||
with set_paths(cwd=b, home=self.tdir):
|
||||
tf(['r', 'd'], {b/'r': '~/b/r', b/'d': '~/b/d', b/'d'/'r': '~/b/d/r'})
|
||||
opts.mode = 'normal'
|
||||
with set_paths(cwd='/some/else', home='/foo/bar'):
|
||||
tf([b/'r', b/'d', '/dest'], {b/'r': '/dest/r', b/'d': '/dest/d', b/'d'/'r': '/dest/d/r'})
|
||||
tf([b/'r', b/'d', '~/dest'], {b/'r': '~/dest/r', b/'d': '~/dest/d', b/'d'/'r': '~/dest/d/r'})
|
||||
with set_paths(cwd=b, home='/foo/bar'):
|
||||
tf(['r', 'd', '/dest'], {b/'r': '/dest/r', b/'d': '/dest/d', b/'d'/'r': '/dest/d/r'})
|
||||
|
||||
os.symlink('/foo/b', b / 'e')
|
||||
os.symlink('r', b / 's')
|
||||
os.link(b / 'r', b / 'h')
|
||||
with set_paths(cwd='/some/else', home='/foo/bar'):
|
||||
files = gm(b / 'e', 'dest')
|
||||
self.ae(files[0].symbolic_link_target, 'path:/foo/b')
|
||||
files = gm(b / 's', b / 'r', 'dest')
|
||||
self.ae(files[0].symbolic_link_target, 'fid:2')
|
||||
files = gm(b / 'h', 'dest')
|
||||
self.ae(files[0].file_type, FileType.regular)
|
||||
files = gm(b / 'h', b / 'r', 'dest')
|
||||
self.ae(files[1].file_type, FileType.link)
|
||||
self.ae(files[1].hard_link_target, '1')
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user