Source code for omni_json_db.utils

# pylint: disable=too-many-lines
from __future__ import annotations
from collections import defaultdict
from contextlib import contextmanager
from abc import ABCMeta
from threading import Lock, Event, Condition, get_ident
from signal import SIGINT, signal, default_int_handler # SIG_IGN
from typing import Callable, Any
#-----------------------------------------------------------------------------
try:
    import ipdb
    debug_break = ipdb.set_trace

except ImportError:
    debug_break = breakpoint

[docs] class JError(Exception): pass
[docs] class JKeyError(JError, KeyError): pass
[docs] class JValueError(JError, ValueError): pass
[docs] class JTypeError(JError, TypeError): pass
#----------------------------------------------------------------------------- # pylint: disable=too-few-public-methods class JDbBase(metaclass=ABCMeta): # pragma: no cover pass # pylint: disable=too-few-public-methods class JIoBase(metaclass=ABCMeta): # pragma: no cover pass
[docs] def deepcopy(src:Any) -> Any: """Create a selective deep copy optimised for the types used in JDb. Common immutable types and :class:`JDbBase` instances are returned as-is without copying. Containers are handled as follows: * ``tuple`` – new tuple whose elements are recursively deep-copied. * ``dict`` – new dict whose *values* are recursively deep-copied (keys are not copied because dict keys must be hashable). * ``set`` – shallow copy via ``set.copy()`` (set elements are hashable scalars and need no further copying). * Any other object whose ``__hash__`` attribute is truthy (e.g. a compiled :class:`re.Pattern` or a ``frozenset``) is treated as effectively immutable and returned without copying. * Everything else (typically a ``list``) – new list whose elements are recursively deep-copied. Args: src (Any): The object to copy. Returns: Any: A deep copy of *src*, or *src* itself for immutable types. Example: >>> original = {'key': [1, 2, 3]} >>> copied = deepcopy(original) >>> copied['key'] is original['key'] False """ if src is None or isinstance(src, (str, bytes, int, float, bool, JDbBase)): return src if isinstance(src, tuple): return tuple(deepcopy(v) for v in src) if isinstance(src, dict): return {key:deepcopy(val) for key, val in src.items()} if isinstance(src, set): return src.copy() if src.__hash__: return src return [deepcopy(val) for val in src]
#----------------------------------------------------------------------------- def Style(msg, bold=None, dim=None, smso=None, underscore=None, blink=None, reverse=None, hidden=None, bright=None, fg=None, black=None, red=None, green=None, yellow=None, blue=None, magenta=None, cyan=None, white=None, bg=None, bg_black=None, bg_red=None, bg_green=None, bg_yellow=None, bg_blue=None, bg_magenta=None, bg_cyan=None, bg_white=None): """Wrap a string in ANSI escape codes to apply terminal colour and text styling. If no styling flags are set, *msg* is returned unchanged. All boolean parameters default to ``None`` (off). **Foreground colour precedence** – ``fg`` overrides the named colour shortcuts (``black``, ``red``, … ``white``). Only the *first* truthy shortcut is applied. **Colour encoding for** ``fg`` **and** ``bg``: * ``int`` (0–7) – standard ANSI colour index directly. * ``str`` – bit-mapped from the characters present: ``'r'`` → +1, ``'g'`` → +2, ``'b'`` → +4. E.g. ``'rg'`` → yellow (3). * ``tuple`` / ``list`` – three-element sequence ``[r, g, b]`` where each value is 0 or 1, bit-mapped the same way. When ``bright=True`` the foreground uses high-intensity ANSI codes (90–97) instead of standard codes (30–37). Args: msg (str): The text to style. bold (bool, optional): Bold / increased intensity. dim (bool, optional): Dim / decreased intensity. smso (bool, optional): Standout mode (terminal-defined highlight). underscore (bool, optional): Underline the text. blink (bool, optional): Blinking text. reverse (bool, optional): Swap foreground and background colours. hidden (bool, optional): Hide the text (invisible). bright (bool, optional): Use high-intensity foreground colour codes. fg (int | str | tuple | list, optional): Foreground colour; see colour encoding above. black (bool, optional): Set foreground colour to black. red (bool, optional): Set foreground colour to red. green (bool, optional): Set foreground colour to green. yellow (bool, optional): Set foreground colour to yellow. blue (bool, optional): Set foreground colour to blue. magenta (bool, optional): Set foreground colour to magenta. cyan (bool, optional): Set foreground colour to cyan. white (bool, optional): Set foreground colour to white. bg (int | str | tuple | list, optional): Background colour; see colour encoding above. bg_black (bool, optional): Set background colour to black. bg_red (bool, optional): Set background colour to red. bg_green (bool, optional): Set background colour to green. bg_yellow (bool, optional): Set background colour to yellow. bg_blue (bool, optional): Set background colour to blue. bg_magenta (bool, optional): Set background colour to magenta. bg_cyan (bool, optional): Set background colour to cyan. bg_white (bool, optional): Set background colour to white. Returns: str: *msg* wrapped in ANSI escape codes, or *msg* unchanged if no styling is requested. Example: >>> print(Style("OK", green=True, bold=True)) >>> print(Style("ERROR", fg='r', bold=True)) >>> print(Style("INFO", fg=[0, 0, 1], bg=0)) # blue on black """ code = '' for ii,vv in enumerate([bold, dim, smso, underscore, blink, reverse, hidden]): if not vv: continue code += f'\033[{ii+1}m' if fg is None: for ii,vv in enumerate([black, red, green, yellow, blue, magenta, cyan, white]): if not vv: continue v1 = 9 if bool(bright) else 3 code += f'\033[{v1}{ii}m' break else: if isinstance(fg, int): vv = max(min(fg, 7), 0) elif isinstance(fg, str): vv = 1 * ('r' in fg) + 2 * ('g' in fg) + 4 * ('b' in fg) else: vv = 1 * fg[0] + 2 * fg[1] + 4 * fg[2] v1 = 9 if bool(bright) else 3 code += f'\033[{v1}{vv}m' if bg is None: for ii,vv in enumerate([bg_black, bg_red, bg_green, bg_yellow, bg_blue, bg_magenta, bg_cyan, bg_white]): if not vv: continue code += f'\033[4{ii}m' break else: if isinstance(bg, int): vv = max(min(bg, 7), 0) elif isinstance(bg, str): vv = 1 * ('r' in bg) + 2 * ('g' in bg) + 4 * ('b' in bg) else: vv = 1 * bg[0] + 2 * bg[1] + 4 * bg[2] code += f'\033[4{vv}m' if not code: return msg return f'{code}{msg}\033[0m' #--------------------------------------------------------------------- #--------------------------------------------------------------------- #--------------------------------------------------------------------- #--------------------------------------------------------------------- class INT_Handler: """Deferred SIGINT handler that protects critical sections from keyboard interrupts. When code enters a protected section (via :meth:`disable`) any ``Ctrl+C`` (SIGINT) is captured silently and recorded instead of raising :exc:`KeyboardInterrupt` immediately. Once all protected sections have exited (via :meth:`enable`), callers can check :meth:`is_called` to discover whether an interrupt was received and act accordingly. This is used internally by :class:`FileLock` to prevent SIGINT from interrupting a write-locked database operation mid-transaction. """ __slots__ = ('count', 'lock', 'call_flag') def __init__(self): """Set up the deferred SIGINT handler and install it as the process SIGINT handler.""" self.count = 0 self.count = 0 self.lock = Lock() self.call_flag = Event() signal(SIGINT, self.handler) def disable(self): """Enter a protected section where SIGINT is deferred rather than raised. Increments the internal nesting counter. If this is the outermost ``disable()`` call (counter was 0), the pending-interrupt flag is cleared so stale events from a previous section cannot bleed through. This method is re-entrant: multiple nested calls are allowed and each must be matched by a corresponding :meth:`enable` call. """ with self.lock: count = self.count self.count = count + 1 if count == 0: self.call_flag.clear() def enable(self): """Leave a protected section, decrementing the nesting counter. When the counter reaches zero the pending-interrupt flag is cleared, discarding any deferred SIGINT that was recorded during the section. Callers should check :meth:`is_called` *before* calling ``enable()`` if they need to act on a deferred interrupt. The counter is never decremented below zero. """ with self.lock: count = self.count = max(0, self.count-1) if count == 0: self.call_flag.clear() def reset(self): """Force-reset the nesting counter to zero and clear the pending-interrupt flag. Use this only in emergency cleanup paths (e.g. after an unhandled exception) where normal :meth:`enable` pairing is not possible. """ with self.lock: # pragma: no cover self.count = 0 self.call_flag.clear() def is_called(self) -> bool: """Return whether a SIGINT was received while inside a protected section. Returns ``True`` only if the pending-interrupt flag is set *and* the nesting counter is still greater than zero (i.e. the signal arrived inside an active protected section that has not yet been fully exited). Returns: bool: ``True`` if a deferred interrupt is pending, ``False`` otherwise. """ if self.call_flag.is_set(): with self.lock: # pragma: no cover return self.count > 0 and self.call_flag.is_set() return False def handler(self, signum, frame): #pragma: no cover """SIGINT signal handler installed at construction time. If no protected section is active (``count == 0``), the default interrupt handler is invoked immediately, which raises :exc:`KeyboardInterrupt` in the normal way. If a protected section is active (``count > 0``), the signal is captured silently and recorded via the pending-interrupt flag so that :meth:`is_called` returns ``True`` after the section exits. Args: signum (int): Signal number (always ``signal.SIGINT`` here). frame (frame): Current stack frame at the point the signal arrived. """ with self.lock: count = self.count if count == 0: self.call_flag.clear() default_int_handler(signum, frame) else: self.call_flag.set() INT_manager = INT_Handler() #--------------------------------------------------------------------- #--------------------------------------------------------------------- #--------------------------------------------------------------------- #---------------------------------------------------------------------
[docs] class FileLockException(BlockingIOError): """Raised when a :class:`FileLock` operation cannot be completed. Thrown in two situations: * A non-blocking lock acquisition (``block=False``) fails because another process already holds an incompatible lock. * A lock acquisition is attempted after the :class:`FileLock` has been closed or is being destroyed (mode ``'x'``). """ pass # pylint: disable=unnecessary-pass
#--------------------------------------------------------------------- #--------------------------------------------------------------------- #--------------------------------------------------------------------- #---------------------------------------------------------------------
[docs] class FileLock: """Combined thread-level and process-level read/write lock backed by OS file locks. Wraps a set of OS-level file-lock callables to provide: * **Shared read locks** (``mode='r'``) – multiple threads *and* processes may hold a read lock simultaneously. * **Exclusive write locks** (``mode='w'``) – only one thread in one process may hold a write lock; all readers are excluded. * **Re-entrant acquisition** – the same thread may call :meth:`acquire` multiple times; each call must be matched by a :meth:`release` call. * **Lock upgrade** – a thread holding a read lock may promote it to a write lock without fully releasing via ``switch=True`` in :meth:`acquire`. * **SIGINT protection** – write locks automatically engage :class:`INT_Handler` so that ``Ctrl+C`` is deferred until the write section completes. Internal mode values stored in ``_mode``: * ``''`` – no lock held. * ``'r'`` – shared read lock active. * ``'w'`` – exclusive write lock active. * ``'p'`` – pending: a thread is waiting for the OS-level lock. * ``'x'`` – closed/destroyed; no new acquisitions are permitted. """ __slots__ = ('_rlock', '_wlock', '_unlock', '_close', '_remove', \ '_lock', '_cond', '_idents', '_mode', 'SIGINT')
[docs] def __init__(self, \ rlock:Callable[[bool], None], wlock:Callable[[bool], None], unlock:Callable[[], None], close:Callable[[], None], remove:Callable[[], None]): """Initialise the lock with OS-level locking callables. Args: rlock (Callable[[bool], None]): Acquire a shared (read) OS-level file lock. The single ``bool`` argument indicates whether the call should block. wlock (Callable[[bool], None]): Acquire an exclusive (write) OS-level file lock. The single ``bool`` argument indicates whether the call should block. unlock (Callable[[], None]): Release the current OS-level file lock. close (Callable[[], None]): Close the underlying lock file handle. remove (Callable[[], None]): Delete the lock file from disk. Raises: TypeError: If any of the five arguments is not callable. """ if not callable(rlock) or not callable(wlock) or not callable(unlock) or not callable(close) or not callable(remove): raise TypeError self._rlock = rlock self._wlock = wlock self._unlock = unlock self._close = close self._remove = remove self._lock = Lock() self._cond = Condition(self._lock) self._idents = defaultdict(int) self._mode = '' self.SIGINT = INT_manager
[docs] def __repr__(self) -> str: """Return a diagnostic string showing the lock's current state. Returns: str: A string of the form ``<FileLock lock:<bool> mode:<mode> at <hex_address>>``. ``lock`` is ``1`` when a read or write lock is active, ``0`` otherwise; ``mode`` is one of ``''``, ``'r'``, ``'w'``, ``'p'``, or ``'x'``. """ return f'<{type(self).__name__} lock:{int(self.is_locked)} mode:{self._mode} at {hex(id(self))}>'
[docs] def __del__(self): """Clean up on garbage collection: release all pending locks and close the lock file.""" self.release_all() self._close()
[docs] def release_all(self) -> bool: # pragma: no cover """Wait for all threads to release their locks, then mark the lock as destroyed. Blocks until ``_idents`` is empty (every thread has called :meth:`release` enough times to drop its count to zero). Once drained, the mode is set to ``'x'`` to prevent any new :meth:`acquire` calls from succeeding. If a write lock was active, the SIGINT handler is re-enabled before closing. This method is called by :meth:`__del__` and should not normally be called directly. Returns: bool: ``True`` if the internal mutex was acquired and the shutdown sequence completed. ``False`` if the mutex itself could not be acquired (should not happen in practice). """ if not self._lock.acquire(): # pylint: disable=consider-using-with return False try: while self._idents: self._cond.wait() if self._mode == 'w': self.SIGINT.enable() self._mode = 'x' self._idents.clear() self._cond.notify_all() finally: self._lock.release() return True
[docs] def reset_lock(self) -> None: # pragma: no cover """Delete the lock file from disk, ignoring the error if it does not exist. Use this to clean up a stale lock file left behind by a crashed process. Only call this when no other process holds or awaits the lock. """ try: self._remove() except FileNotFoundError: pass
@property def is_locked(self) -> bool: """Whether any thread currently holds a read or write lock. Returns: bool: ``True`` if ``mode`` is ``'r'`` or ``'w'``, ``False`` otherwise. """ return self._mode == 'r' or self._mode == 'w' @property def mode(self) -> str: """Current lock mode as a single character string. Returns: str: One of: * ``''`` – no lock held. * ``'r'`` – shared read lock active. * ``'w'`` – exclusive write lock active. * ``'p'`` – a thread is blocked waiting for the OS-level lock. * ``'x'`` – lock is closed; no new acquisitions permitted. """ return self._mode
[docs] @contextmanager def rlock(self): # pragma: no cover """Context manager that acquires a shared read lock and releases it on exit. Yields: None: Control is yielded to the ``with`` block with the read lock held. Example: :: with file_lock.rlock(): data = read_from_file() """ self.acquire(read_only=True) try: yield finally: self.release()
[docs] @contextmanager def wlock(self): # pragma: no cover """Context manager that acquires an exclusive write lock and releases it on exit. SIGINT (``Ctrl+C``) is deferred while the write lock is held and re-enabled automatically on release. Yields: None: Control is yielded to the ``with`` block with the write lock held. Example: :: with file_lock.wlock(): write_to_file(data) """ self.acquire(read_only=False) try: yield finally: self.release()
[docs] def has_SIGINT(self) -> bool: """Return whether a ``Ctrl+C`` was received while a write lock was held. This delegates to :meth:`INT_Handler.is_called` on the shared :data:`INT_manager` instance. Returns: bool: ``True`` if a deferred SIGINT is pending, ``False`` otherwise. """ return self.SIGINT.is_called()
[docs] def can_lock(self) -> bool: """Test whether an exclusive write lock can be acquired immediately without blocking. Attempts a non-blocking ``acquire(block=False, read_only=False)`` and releases it straight away. Returns: bool: ``True`` if the write lock was obtained (and released), ``False`` if another holder would have caused a block. """ try: self.acquire(block=False, read_only=False) return True except FileLockException: # pragma: no cover return False finally: self.release()
[docs] def get_count(self, thread_id:int) -> int: """Return the re-entrance count for a given thread. Each call to :meth:`acquire` increments the count for the calling thread; each :meth:`release` decrements it. The OS-level lock is released only when the count returns to zero. Args: thread_id (int): Thread identifier as returned by :func:`threading.get_ident`. Returns: int: Number of times the thread has acquired this lock without a matching release. Returns ``0`` if the thread holds no lock. """ return self._idents.get(thread_id, 0)
[docs] def acquire(self, block:bool=True, read_only:bool=False, switch:bool=False) -> int: """Acquire a read or write lock for the calling thread. Thread-level re-entrance is supported: calling ``acquire`` again from a thread that already holds a compatible lock simply increments the re-entrance counter and returns immediately. **Lock promotion (** ``switch=True`` **)** – a thread that currently holds a read lock may atomically promote it to a write lock. The read lock is released and the write lock is acquired without allowing other threads to sneak in between. Args: block (bool, optional): If ``True`` (default), block until the lock becomes available. If ``False``, raise :exc:`FileLockException` immediately when the lock cannot be acquired. read_only (bool, optional): If ``True``, acquire a shared read lock (multiple threads/processes may hold it simultaneously). If ``False`` (default), acquire an exclusive write lock. switch (bool, optional): If ``True``, upgrade the current thread's read lock to a write lock without fully releasing. Only valid when the calling thread already holds a read lock. Defaults to ``False``. Returns: int: The calling thread's identifier (as returned by :func:`threading.get_ident`). Raises: RuntimeError: If the internal threading mutex cannot be acquired. FileLockException: If ``block=False`` and the lock is held by another thread or process, or if the lock has been closed (mode ``'x'``). """ if not self._lock.acquire(): # pylint: disable=consider-using-with raise RuntimeError try: ident = get_ident() _idents = self._idents while True: _mode = self._mode if _mode == 'x': # pragma: no cover raise FileLockException("FileLock is closed or being destroyed.") # [1] Thread level if _mode == 'r' and read_only and _idents: # allow multiple reader _idents[ident] += 1 return ident if _mode == 'w' and ident in _idents: # only allow one writer (same thread) _idents[ident] += 1 return ident elif _mode == 'r' and ident in _idents: # switch 'r' to 'w' _cnt = _idents[ident] if _cnt <= 1: _idents.pop(ident) if not _idents: # this thread is the last lock owner try: self._unlock() except OSError as e: # pragma: no cover print(e) if self._mode == 'x': continue _mode = self._mode = '' if not switch: self._cond.notify_all() _idents[ident] = _cnt continue if not switch and _cnt > 0: # pragma: no cover _idents[ident] = _cnt elif switch: # pragma: no cover _idents[ident] = _cnt - 1 if _cnt > 1 and len(_idents) == 1: try: self._unlock() except OSError as e: print(e) if self._mode == 'x': continue _mode = self._mode = '' if _mode != '': # pragma: no cover if not block: raise FileLockException("Could not acquire lock") # pragma: no cover self._cond.wait() continue # [2] process level try: if read_only: self._rlock(block=False) self._mode = 'r' else: self._wlock(block=False) self._mode = 'w' self.SIGINT.disable() _idents[ident] += 1 self._cond.notify_all() return ident except BlockingIOError as e: if not block: # pragma: no cover if ident in _idents: self._mode = 'r' raise FileLockException("Could not acquire lock") from e self._mode = 'p' self._lock.release() os_lock_acquired = False os_err = None if self._mode != 'p': # pragma: no cover continue try: if read_only: self._rlock(block=True) else: self._wlock(block=True) os_lock_acquired = True except Exception as ex: # pragma: no cover os_err = ex finally: self._lock.acquire() # pylint: disable=consider-using-with if self._mode == 'p': # pragma: no cover self._mode = '' self._cond.notify_all() if self._mode == 'x': # pragma: no cover if os_lock_acquired: try: self._unlock() except OSError as e1: print(e1) raise FileLockException("FileLock is closed or being destroyed.") from e if os_err is not None: # pragma: no cover raise FileLockException("Could not acquire lock") from os_err if read_only: self._mode = 'r' else: self._mode = 'w' self.SIGINT.disable() _idents[ident] += 1 self._cond.notify_all() # wake up all thread due to 'p' return ident finally: self._lock.release() return ident
[docs] def release(self) -> int: """Release one acquisition of the lock for the calling thread. Decrements the re-entrance counter for the calling thread. When the counter reaches zero and no other threads hold the lock, the OS-level file lock is released and SIGINT handling is re-enabled (if a write lock was held). Calling ``release`` from a thread that does not hold the lock has no effect. Returns: int: The calling thread's identifier (as returned by :func:`threading.get_ident`). Raises: RuntimeError: If the internal threading mutex cannot be acquired. """ if not self._lock.acquire(): # pylint: disable=consider-using-with raise RuntimeError try: _idents = self._idents ident = get_ident() if ident in _idents: if _idents.get(ident, 0) <= 1: _idents.pop(ident, 0) else: _idents[ident] -= 1 if not _idents: if self._mode == 'w': self.SIGINT.enable() try: self._unlock() except OSError as e1: # pragma: no cover print(e1) self._mode = '' self._cond.notify_all() return ident finally: self._lock.release()
#