145 lines
4.5 KiB
Python
145 lines
4.5 KiB
Python
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)
|