DRYer

kittens/transfer/send.py

@@ -18,7 +18,7 @@ from kitty.cli_stub import TransferCLIOptions
 from kitty.fast_data_types import FILE_TRANSFER_CODE, wcswidth
 from kitty.file_transmission import (
     Action, Compression, FileTransmissionCommand, FileType, NameReprEnum,
-    TransmissionType, encode_bypass
+    TransmissionType, encode_bypass, split_for_transfer
 )
 from kitty.typing import KeyEventType
 from kitty.utils import sanitize_control_codes
@@ -342,15 +342,9 @@ class SendManager:
             chunk, usz = af.next_chunk()
             self.current_chunk_uncompressed_sz += usz
         is_last = af.state is FileState.finished
-        mv = memoryview(chunk)
-        pos = 0
-        limit = len(chunk)
-        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()
+        if len(chunk):
+            for ftc in split_for_transfer(chunk, file_id=af.file_id, mark_last=is_last):
+                yield ftc.serialize()
         elif is_last:
             yield FileTransmissionCommand(action=Action.end_data, file_id=af.file_id, data=b'').serialize()
 

kitty/file_transmission.py

@@ -46,6 +46,22 @@ def encode_bypass(request_id: str, bypass: str) -> str:
     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):
 
     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))
             return
         has_capacity = True
-        while len(chunk) > 0:
-            data = FileTransmissionCommand(id=receive_id, action=Action.data, file_id=file_id, data=chunk[:4096])
+        for data in split_for_transfer(chunk, session_id=receive_id, file_id=file_id):
             if has_capacity:
                 if not self.write_ftc_to_child(data, use_pending=False):
                     has_capacity = False
                     pending.append(data)
             else:
                 pending.append(data)
-            chunk = chunk[4096:]
         self.callback_after(func)
 
     def send_status_response(