diff --git a/kittens/transfer/librsync.py b/kittens/transfer/librsync.py index 72616176b..4a3f6674f 100644 --- a/kittens/transfer/librsync.py +++ b/kittens/transfer/librsync.py @@ -3,34 +3,39 @@ # License: GPLv3 Copyright: 2021, Kovid Goyal import os -from typing import Generator, Iterator, Optional +from typing import IO, Generator, Iterator, Optional from .rsync import ( - IO_BUFFER_SIZE, RsyncError, SignatureCapsule, begin_create_signature, - begin_load_signature, iter_job, make_hash_table + IO_BUFFER_SIZE, JobCapsule, RsyncError, SignatureCapsule, + begin_create_delta, begin_create_signature, begin_load_signature, iter_job, + make_hash_table ) +def drive_job_on_file(f: IO[bytes], job: JobCapsule) -> Iterator[bytes]: + finished = False + prev_unused_input = b'' + while not finished: + input_data = f.read(IO_BUFFER_SIZE) + no_more_data = not input_data + if prev_unused_input: + input_data = prev_unused_input + input_data + prev_unused_input = b'' + output, finished, sz_of_unused_input = iter_job(job, input_data, no_more_data) + if sz_of_unused_input > 0 and not finished: + if no_more_data: + raise RsyncError(f"{sz_of_unused_input} bytes of input data were not used") + prev_unused_input = input_data[-sz_of_unused_input:] + yield output + + def signature_of_file(path: str) -> Iterator[bytes]: with open(path, 'rb') as f: f.seek(0, os.SEEK_END) fsz = f.tell() job = begin_create_signature(fsz) f.seek(0) - finished = False - prev_unused_input = b'' - while not finished: - input_data = f.read(IO_BUFFER_SIZE) - no_more_data = not input_data - if prev_unused_input: - input_data = prev_unused_input + input_data - prev_unused_input = b'' - output, finished, sz_of_unused_input = iter_job(job, input_data, no_more_data) - if sz_of_unused_input > 0 and not finished: - if no_more_data: - raise RsyncError(f"{sz_of_unused_input} bytes of input data were not used") - prev_unused_input = input_data[-sz_of_unused_input:] - yield output + yield from drive_job_on_file(f, job) def load_signature() -> Generator[Optional[SignatureCapsule], bytes, None]: @@ -50,3 +55,9 @@ def load_signature() -> Generator[Optional[SignatureCapsule], bytes, None]: prev_unused_input = input_data[-sz_of_unused_input:] make_hash_table(signature) yield signature + + +def delta_for_file(path: str, sig: SignatureCapsule) -> Iterator[bytes]: + job = begin_create_delta(sig) + with open(path, 'rb') as f: + yield from drive_job_on_file(f, job) diff --git a/kittens/transfer/rsync.c b/kittens/transfer/rsync.c index 57187ccae..06d56b339 100644 --- a/kittens/transfer/rsync.c +++ b/kittens/transfer/rsync.c @@ -123,7 +123,7 @@ build_hash_table(PyObject *self UNUSED, PyObject *args) { static PyObject* begin_create_delta(PyObject *self UNUSED, PyObject *args) { PyObject *sig_capsule; - if (!PyArg_ParseTuple(args, "O!y#|p", &PyCapsule_Type, &sig_capsule)) return NULL; + if (!PyArg_ParseTuple(args, "O!", &PyCapsule_Type, &sig_capsule)) return NULL; GET_SIG_FROM_CAPSULE; rs_job_t *job = rs_delta_begin(sig); if (!job) return PyErr_NoMemory(); diff --git a/kittens/transfer/rsync.pyi b/kittens/transfer/rsync.pyi index b1b2b5156..7a8572848 100644 --- a/kittens/transfer/rsync.pyi +++ b/kittens/transfer/rsync.pyi @@ -27,5 +27,9 @@ def make_hash_table(sig: SignatureCapsule) -> None: pass +def begin_create_delta(sig: SignatureCapsule) -> JobCapsule: + pass + + def iter_job(job_capsule: JobCapsule, input_data: bytes, eof: Optional[bool] = None, expecting_output: bool = True) -> Tuple[bytes, bool, int]: pass