95 lines
2.6 KiB
Python
95 lines
2.6 KiB
Python
from __future__ import annotations
|
|
|
|
import builtins
|
|
import io
|
|
import os
|
|
import threading
|
|
|
|
from ._api import _ORIGINAL_OPEN, _validate_path_mode, _wrap_reader, open as dlp_open
|
|
|
|
_ORIGINAL_IO_OPEN = io.open
|
|
_patch_lock = threading.RLock()
|
|
_installed = False
|
|
|
|
|
|
def _reader_wrapper(open_read_func):
|
|
def patched_open(
|
|
file,
|
|
mode="r",
|
|
buffering=-1,
|
|
encoding=None,
|
|
errors=None,
|
|
newline=None,
|
|
closefd=True,
|
|
opener=None,
|
|
):
|
|
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")
|
|
reader = open_read_func(os.fsdecode(os.fspath(file)))
|
|
if isinstance(reader, io.RawIOBase):
|
|
return _wrap_reader(reader, mode, buffering, encoding, errors, newline)
|
|
if "b" in mode:
|
|
return reader
|
|
return io.TextIOWrapper(
|
|
reader,
|
|
encoding=encoding,
|
|
errors=errors,
|
|
newline=newline,
|
|
line_buffering=buffering == 1,
|
|
)
|
|
return _ORIGINAL_OPEN(file, mode, buffering, encoding, errors, newline, closefd)
|
|
|
|
return patched_open
|
|
|
|
|
|
def install_open_patch() -> None:
|
|
"""Route global path-based reads through the lazy default dlp-io session."""
|
|
global _installed
|
|
with _patch_lock:
|
|
if _installed:
|
|
return
|
|
builtins.open = dlp_open
|
|
io.open = dlp_open
|
|
_installed = True
|
|
|
|
|
|
def install_reader_patch(open_read_func) -> None:
|
|
"""Compatibility hook for an explicitly managed session.open_raw callback."""
|
|
global _installed
|
|
with _patch_lock:
|
|
if _installed:
|
|
return
|
|
wrapper = _reader_wrapper(open_read_func)
|
|
builtins.open = wrapper
|
|
io.open = wrapper
|
|
_installed = True
|
|
|
|
|
|
def uninstall_open_patch() -> None:
|
|
global _installed
|
|
with _patch_lock:
|
|
if not _installed:
|
|
return
|
|
builtins.open = _ORIGINAL_OPEN
|
|
io.open = _ORIGINAL_IO_OPEN
|
|
_installed = False
|