Initial commit: split dlp-io library out of DataPacker

2026-07-31 14:58:37 +08:00
commit 3b0c81fd9d
22 changed files with 2509 additions and 0 deletions
@@ -0,0 +1,134 @@
+from __future__ import annotations
+
+import io
+
+from ._errors import DlpTransportError
+from ._helper import CHUNK_SIZE
+
+
+class _PipeRawIO(io.RawIOBase):
+    _DISCARD_LIMIT = 64 * 1024 * 1024
+
+    def __init__(self, *, conn, size, path, offset, reopen, release) -> None:
+        super().__init__()
+        self._conn = conn
+        self._size = int(size)
+        self._path = path
+        self._pos = int(offset)
+        self._expected_end = max(self._size, self._pos)
+        self._reopen = reopen
+        self._release = release
+        self._buffer = bytearray()
+        self._eof = False
+        self._released = False
+
+    def readable(self) -> bool:
+        return True
+
+    def seekable(self) -> bool:
+        return True
+
+    def tell(self) -> int:
+        if self.closed:
+            raise ValueError("I/O operation on closed file")
+        return self._pos
+
+    def _receive(self) -> None:
+        try:
+            chunk = self._conn.recv_bytes(maxlength=CHUNK_SIZE)
+        except (EOFError, OSError) as exc:
+            raise DlpTransportError("helper disconnected before EOF") from exc
+        if chunk:
+            received_end = self._pos + len(self._buffer) + len(chunk)
+            if received_end > self._expected_end:
+                raise DlpTransportError(
+                    "file size changed during transfer: expected end %d, received at least %d"
+                    % (self._expected_end, received_end)
+                )
+            self._buffer.extend(chunk)
+            return
+        if self._pos != self._expected_end:
+            raise DlpTransportError(
+                "file size changed during transfer: expected end %d, received %d"
+                % (self._expected_end, self._pos)
+            )
+        self._eof = True
+
+    def readinto(self, target) -> int:
+        if self.closed:
+            raise ValueError("I/O operation on closed file")
+        if len(target) == 0:
+            return 0
+        if not self._buffer:
+            if self._eof:
+                return 0
+            self._receive()
+            if self._eof:
+                return 0
+        count = min(len(target), len(self._buffer))
+        target[:count] = self._buffer[:count]
+        del self._buffer[:count]
+        self._pos += count
+        return count
+
+    def seek(self, offset, whence=io.SEEK_SET) -> int:
+        if self.closed:
+            raise ValueError("I/O operation on closed file")
+        if whence == io.SEEK_SET:
+            target = offset
+        elif whence == io.SEEK_CUR:
+            target = self._pos + offset
+        elif whence == io.SEEK_END:
+            target = self._size + offset
+        else:
+            raise ValueError("invalid whence: %r" % (whence,))
+        if target < 0:
+            raise ValueError("negative seek position: %d" % target)
+        delta = target - self._pos
+        if delta == 0:
+            return self._pos
+        if 0 < delta <= self._DISCARD_LIMIT:
+            remaining = delta
+            while remaining:
+                scratch = bytearray(min(remaining, CHUNK_SIZE))
+                count = self.readinto(scratch)
+                if count == 0:
+                    self._pos = target
+                    self._expected_end = max(self._size, target)
+                    break
+                remaining -= count
+            return self._pos
+
+        self._conn.close()
+        try:
+            conn, size = self._reopen(self._path, target)
+        except Exception:
+            self.close()
+            raise
+        if int(size) != self._size:
+            conn.close()
+            self.close()
+            raise DlpTransportError(
+                "file size changed during seek: initial %d, reopened %d"
+                % (self._size, size)
+            )
+        self._conn = conn
+        self._pos = target
+        self._expected_end = max(self._size, target)
+        self._buffer.clear()
+        self._eof = False
+        return self._pos
+
+    def _release_once(self) -> None:
+        if self._released:
+            return
+        self._released = True
+        self._release(self)
+
+    def close(self) -> None:
+        if not self.closed:
+            try:
+                self._conn.close()
+            finally:
+                self._release_once()
+        super().close()