DRYer
This commit is contained in:
parent
14f9f64bce
commit
f6e518d2f9
@ -18,7 +18,7 @@ from kitty.cli_stub import TransferCLIOptions
|
|||||||
from kitty.fast_data_types import FILE_TRANSFER_CODE, wcswidth
|
from kitty.fast_data_types import FILE_TRANSFER_CODE, wcswidth
|
||||||
from kitty.file_transmission import (
|
from kitty.file_transmission import (
|
||||||
Action, Compression, FileTransmissionCommand, FileType, NameReprEnum,
|
Action, Compression, FileTransmissionCommand, FileType, NameReprEnum,
|
||||||
TransmissionType, encode_bypass
|
TransmissionType, encode_bypass, split_for_transfer
|
||||||
)
|
)
|
||||||
from kitty.typing import KeyEventType
|
from kitty.typing import KeyEventType
|
||||||
from kitty.utils import sanitize_control_codes
|
from kitty.utils import sanitize_control_codes
|
||||||
@ -342,15 +342,9 @@ class SendManager:
|
|||||||
chunk, usz = af.next_chunk()
|
chunk, usz = af.next_chunk()
|
||||||
self.current_chunk_uncompressed_sz += usz
|
self.current_chunk_uncompressed_sz += usz
|
||||||
is_last = af.state is FileState.finished
|
is_last = af.state is FileState.finished
|
||||||
mv = memoryview(chunk)
|
if len(chunk):
|
||||||
pos = 0
|
for ftc in split_for_transfer(chunk, file_id=af.file_id, mark_last=is_last):
|
||||||
limit = len(chunk)
|
yield ftc.serialize()
|
||||||
if limit:
|
|
||||||
while pos < limit:
|
|
||||||
cc = mv[pos:pos + 4096]
|
|
||||||
pos += 4096
|
|
||||||
final = is_last and pos >= limit
|
|
||||||
yield FileTransmissionCommand(action=Action.end_data if final else Action.data, file_id=af.file_id, data=cc).serialize()
|
|
||||||
elif is_last:
|
elif is_last:
|
||||||
yield FileTransmissionCommand(action=Action.end_data, file_id=af.file_id, data=b'').serialize()
|
yield FileTransmissionCommand(action=Action.end_data, file_id=af.file_id, data=b'').serialize()
|
||||||
|
|
||||||
|
|||||||
@ -46,6 +46,22 @@ def encode_bypass(request_id: str, bypass: str) -> str:
|
|||||||
return 'sha256:' + hashlib.sha256(q.encode('utf-8', 'replace')).hexdigest()
|
return 'sha256:' + hashlib.sha256(q.encode('utf-8', 'replace')).hexdigest()
|
||||||
|
|
||||||
|
|
||||||
|
def split_for_transfer(
|
||||||
|
data: Union[bytes, bytearray, memoryview],
|
||||||
|
session_id: str = '', file_id: str = '',
|
||||||
|
mark_last: bool = False,
|
||||||
|
chunk_size: int = 4096
|
||||||
|
) -> Iterator['FileTransmissionCommand']:
|
||||||
|
if isinstance(data, (bytes, bytearray)):
|
||||||
|
data = memoryview(data)
|
||||||
|
while len(data):
|
||||||
|
ac = Action.data
|
||||||
|
if mark_last and len(data) <= chunk_size:
|
||||||
|
ac = Action.end_data
|
||||||
|
yield FileTransmissionCommand(action=ac, id=session_id, file_id=file_id, data=data[:chunk_size])
|
||||||
|
data = data[chunk_size:]
|
||||||
|
|
||||||
|
|
||||||
class NameReprEnum(Enum):
|
class NameReprEnum(Enum):
|
||||||
|
|
||||||
def __repr__(self) -> str:
|
def __repr__(self) -> str:
|
||||||
@ -604,15 +620,13 @@ class FileTransmission:
|
|||||||
self.write_ftc_to_child(FileTransmissionCommand(id=receive_id, action=Action.end_data, file_id=file_id))
|
self.write_ftc_to_child(FileTransmissionCommand(id=receive_id, action=Action.end_data, file_id=file_id))
|
||||||
return
|
return
|
||||||
has_capacity = True
|
has_capacity = True
|
||||||
while len(chunk) > 0:
|
for data in split_for_transfer(chunk, session_id=receive_id, file_id=file_id):
|
||||||
data = FileTransmissionCommand(id=receive_id, action=Action.data, file_id=file_id, data=chunk[:4096])
|
|
||||||
if has_capacity:
|
if has_capacity:
|
||||||
if not self.write_ftc_to_child(data, use_pending=False):
|
if not self.write_ftc_to_child(data, use_pending=False):
|
||||||
has_capacity = False
|
has_capacity = False
|
||||||
pending.append(data)
|
pending.append(data)
|
||||||
else:
|
else:
|
||||||
pending.append(data)
|
pending.append(data)
|
||||||
chunk = chunk[4096:]
|
|
||||||
self.callback_after(func)
|
self.callback_after(func)
|
||||||
|
|
||||||
def send_status_response(
|
def send_status_response(
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user