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,144 @@
+from __future__ import annotations
+
+import atexit
+import builtins
+import io
+import operator
+import os
+import threading
+
+from ._config import DlpConfig
+from ._errors import DlpConfigurationError
+from ._session import DlpSession
+
+_ORIGINAL_OPEN = builtins.open
+_DEFAULT_BUFFER_SIZE = 1024 * 1024
+_state_lock = threading.RLock()
+_config = DlpConfig()
+_default_session = None
+
+
+def _validate_path_mode(mode, buffering, encoding, errors, newline) -> int:
+    buffering = operator.index(buffering)
+    if buffering < -1:
+        raise ValueError("invalid buffering size")
+    if not isinstance(mode, str):
+        raise TypeError("open() argument 'mode' must be str, not %s" % type(mode).__name__)
+    allowed = set("axrwb+t")
+    invalid = set(mode) - allowed
+    duplicate = any(mode.count(character) > 1 for character in allowed)
+    if not mode or invalid or duplicate or sum(c in mode for c in "axrw") != 1:
+        raise ValueError("invalid mode: %r" % mode)
+    if "b" in mode and "t" in mode:
+        raise ValueError("can't have text and binary mode at once")
+    if "b" in mode:
+        if encoding is not None:
+            raise ValueError("binary mode doesn't take an encoding argument")
+        if errors is not None:
+            raise ValueError("binary mode doesn't take an errors argument")
+        if newline is not None:
+            raise ValueError("binary mode doesn't take a newline argument")
+    elif buffering == 0:
+        raise ValueError("can't have unbuffered text I/O")
+    return buffering
+
+
+def configure(*, python_executable=None, startup_timeout=60) -> None:
+    """Configure the lazy default session before its first bridged read."""
+    global _config
+    with _state_lock:
+        if _default_session is not None:
+            raise DlpConfigurationError(
+                "cannot configure dlp_io while the default session is active"
+            )
+        if startup_timeout <= 0:
+            raise DlpConfigurationError("startup_timeout must be positive")
+        _config = DlpConfig(
+            python_executable=python_executable,
+            startup_timeout=startup_timeout,
+        )
+
+
+def _get_default_session():
+    global _default_session
+    with _state_lock:
+        if _default_session is None:
+            _default_session = DlpSession.start(
+                python_cmd=_config.python_command(),
+                timeout=_config.startup_timeout,
+            )
+        return _default_session
+
+
+def _wrap_reader(raw, mode, buffering, encoding, errors, newline):
+    binary = "b" in mode
+    if binary and buffering == 0:
+        return raw
+    if not binary and buffering == 0:
+        raw.close()
+        raise ValueError("can't have unbuffered text I/O")
+    buffer_size = _DEFAULT_BUFFER_SIZE if buffering in (-1, 1) else buffering
+    if buffer_size <= 0:
+        raw.close()
+        raise ValueError("invalid buffering size")
+    buffered = io.BufferedReader(raw, buffer_size=buffer_size)
+    if binary:
+        return buffered
+    return io.TextIOWrapper(
+        buffered,
+        encoding=encoding,
+        errors=errors,
+        newline=newline,
+        line_buffering=buffering == 1,
+    )
+
+
+def open(
+    file,
+    mode="r",
+    buffering=-1,
+    encoding=None,
+    errors=None,
+    newline=None,
+    closefd=True,
+    opener=None,
+):
+    """Open a file with io.open-compatible calling conventions."""
+    is_path = isinstance(file, (str, bytes, os.PathLike))
+    if not is_path or opener is not None:
+        return _ORIGINAL_OPEN(
+            file,
+            mode,
+            buffering,
+            encoding,
+            errors,
+            newline,
+            closefd,
+            opener,
+        )
+    buffering = _validate_path_mode(mode, buffering, encoding, errors, newline)
+    if "+" in mode:
+        raise io.UnsupportedOperation(
+            "dlp_io does not support update mode for path-based files"
+        )
+    pure_read = "r" in mode
+    if pure_read:
+        if not closefd:
+            raise ValueError("Cannot use closefd=False with file name")
+        path = os.fsdecode(os.fspath(file))
+        raw = _get_default_session().open_raw(path)
+        return _wrap_reader(raw, mode, buffering, encoding, errors, newline)
+    return _ORIGINAL_OPEN(file, mode, buffering, encoding, errors, newline, closefd)
+
+
+def shutdown() -> None:
+    """Close and forget the lazy default session. This operation is idempotent."""
+    global _default_session
+    with _state_lock:
+        session = _default_session
+        _default_session = None
+    if session is not None:
+        session.close()
+
+
+atexit.register(shutdown)