KittyPatch/kitty
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Move management of destination file completely into PatchFile

Browse Source
This commit is contained in:
Kovid Goyal 2021-10-01 14:29:14 +05:30
parent eeb02ceef4
commit e6cff61f99
No known key found for this signature in database
GPG Key ID: 06BC317B515ACE7C
2 changed files with 26 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

26
kittens/transfer/librsync.py
View File

@ -3,11 +3,12 @@
# License: GPLv3 Copyright: 2021, Kovid Goyal <kovid at kovidgoyal.net> # License: GPLv3 Copyright: 2021, Kovid Goyal <kovid at kovidgoyal.net>
import os import os
import tempfile
from typing import IO, TYPE_CHECKING, Iterator from typing import IO, TYPE_CHECKING, Iterator
from .rsync import ( from .rsync import (
IO_BUFFER_SIZE, RsyncError, begin_create_delta, begin_create_signature, IO_BUFFER_SIZE, RsyncError, begin_create_delta, begin_create_signature,
begin_load_signature, begin_patch, iter_job, build_hash_table begin_load_signature, begin_patch, build_hash_table, iter_job
) )
if TYPE_CHECKING: if TYPE_CHECKING:
@ -76,6 +77,7 @@ class LoadSignature(StreamingJob):
# see whole.c in librsync source for size calculations # see whole.c in librsync source for size calculations
expected_input_size = 16 * 1024 expected_input_size = 16 * 1024
autocommit = True
def __init__(self) -> None: def __init__(self) -> None:
job, self.signature = begin_load_signature() job, self.signature = begin_load_signature()
@ -97,8 +99,13 @@ class PatchFile(StreamingJob):
# see whole.c in librsync source for size calculations # see whole.c in librsync source for size calculations
expected_input_size = IO_BUFFER_SIZE expected_input_size = IO_BUFFER_SIZE
def __init__(self, src_path: str): def __init__(self, src_path: str, output_path: str = ''):
self.overwrite_src = not output_path
self.src_file = open(src_path, 'rb') self.src_file = open(src_path, 'rb')
if self.overwrite_src:
self.dest_file = tempfile.NamedTemporaryFile(mode='wb', dir=os.path.dirname(os.path.abspath(os.path.realpath(src_path))), delete=False)
else:
self.dest_file = open(output_path, 'wb')
job = begin_patch(self.read_from_src) job = begin_patch(self.read_from_src)
super().__init__(job, output_buf_size=4 * IO_BUFFER_SIZE) super().__init__(job, output_buf_size=4 * IO_BUFFER_SIZE)
@ -109,7 +116,20 @@ class PatchFile(StreamingJob):
def close(self) -> None: def close(self) -> None:
if not self.src_file.closed: if not self.src_file.closed:
self.src_file.close() self.src_file.close()
commit = close count = 100
while not self.finished:
self()
count -= 1
if count == 0:
raise Exception('Patching file did not receive enough input')
self.dest_file.close()
if self.overwrite_src:
os.replace(self.dest_file.name, self.src_file.name)
def write(self, data: bytes) -> None:
output = self(data)
if output:
self.dest_file.write(output)
def __enter__(self) -> 'PatchFile': def __enter__(self) -> 'PatchFile':
return self return self

11
kitty_tests/file_transmission.py
View File

@ -101,16 +101,11 @@ class TestFileTransmission(BaseTest):
sig_loader(chunk) sig_loader(chunk)
sig_loader() sig_loader()
self.assertTrue(sig_loader.finished) self.assertTrue(sig_loader.finished)
with open(c_path, 'wb') as dest, PatchFile(a_path) as patcher: with PatchFile(a_path, c_path) as patcher:
for chunk in delta_for_file(b_path, sig_loader.signature): for chunk in delta_for_file(b_path, sig_loader.signature):
self.assertFalse(patcher.finished) self.assertFalse(patcher.finished)
output = patcher(chunk) patcher.write(chunk)
if output: self.assertTrue(patcher.finished)
dest.write(output)
while not patcher.finished:
output = patcher()
if output:
dest.write(output)
with open(b_path, 'rb') as b, open(c_path, 'rb') as c: with open(b_path, 'rb') as b, open(c_path, 'rb') as c:
while True: while True:
bc = b.read(4096) bc = b.read(4096)