Source code for portalocker.utils

import os
import time
import tempfile
import contextlib

from . import portalocker

DEFAULT_TIMEOUT = 5
DEFAULT_CHECK_INTERVAL = 0.25
LOCK_METHOD = portalocker.LOCK_EX | portalocker.LOCK_NB

__all__ = [
    'Lock',
    'AlreadyLocked',
    'open_atomic',
]


[docs]class AlreadyLocked(Exception): pass
@contextlib.contextmanager
[docs]def open_atomic(filename, binary=True): '''Open a file for atomic writing. Instead of locking this method allows you to write the entire file and move it to the actual location. Note that is still not atomic in all cases and won't work on existing files. http://docs.python.org/library/os.html#os.rename >>> filename = 'test_file.txt' >>> if os.path.exists(filename): ... os.remove(filename) >>> with open_atomic(filename) as fh: ... fh.write('test') >>> assert os.path.exists(filename) >>> os.remove(filename) ''' assert not os.path.exists(filename), '%r exists' % filename path, name = os.path.split(filename) # Create the parent directory if it doesn't exist if path and not os.path.isdir(path): # pragma: no cover os.makedirs(path) temp_fh = tempfile.NamedTemporaryFile( mode=binary and 'wb' or 'w', dir=path, delete=False, ) yield temp_fh temp_fh.flush() os.fsync(temp_fh.fileno()) temp_fh.close() try: os.rename(temp_fh.name, filename) finally: try: os.remove(temp_fh.name) except Exception: pass
[docs]class Lock(object): def __init__( self, filename, mode='a', truncate=0, timeout=DEFAULT_TIMEOUT, check_interval=DEFAULT_CHECK_INTERVAL, fail_when_locked=True, flags=LOCK_METHOD): '''Lock manager with build-in timeout filename -- filename mode -- the open mode, 'a' or 'ab' should be used for writing truncate -- use truncate to emulate 'w' mode, None is disabled, 0 is truncate to 0 bytes timeout -- timeout when trying to acquire a lock check_interval -- check interval while waiting fail_when_locked -- after the initial lock failed, return an error or lock the file fail_when_locked is useful when multiple threads/processes can race when creating a file. If set to true than the system will wait till the lock was acquired and then return an AlreadyLocked exception. Note that the file is opened first and locked later. So using 'w' as mode will result in truncate _BEFORE_ the lock is checked. ''' self.fh = None self.filename = filename self.mode = mode self.truncate = truncate self.timeout = timeout self.check_interval = check_interval self.fail_when_locked = fail_when_locked self.flags = flags assert 'w' not in mode, 'Mode "w" clears the file before locking'
[docs] def acquire( self, timeout=None, check_interval=None, fail_when_locked=None): '''Acquire the locked filehandle''' if timeout is None: timeout = self.timeout if timeout is None: timeout = 0 if check_interval is None: check_interval = self.check_interval if fail_when_locked is None: fail_when_locked = self.fail_when_locked # If we already have a filehandle, return it fh = self.fh if fh: return fh # Get a new filehandler fh = self._get_fh() try: # Try to lock fh = self._get_lock(fh) except portalocker.LockException as exception: # Try till the timeout is 0 while timeout > 0: # Wait a bit time.sleep(check_interval) timeout -= check_interval # Try again try: # We already tried to the get the lock # If fail_when_locked is true, then stop trying if fail_when_locked: raise AlreadyLocked(exception) else: # pragma: no cover # We've got the lock fh = self._get_lock(fh) break except portalocker.LockException: pass else: # We got a timeout... reraising raise portalocker.LockException(exception) # Prepare the filehandle (truncate if needed) fh = self._prepare_fh(fh) self.fh = fh return fh
[docs] def release(self): '''Releases the currently locked file handle''' if self.fh: self.fh.close() self.fh = None
def _get_fh(self): '''Get a new filehandle''' return open(self.filename, self.mode) def _get_lock(self, fh): ''' Try to lock the given filehandle returns LockException if it fails''' portalocker.lock(fh, self.flags) return fh def _prepare_fh(self, fh, truncate=None): ''' Prepare the filehandle for usage If truncate is a number, the file will be truncated to that amount of bytes ''' if truncate is None: truncate = self.truncate if truncate is not None: fh.seek(truncate) fh.truncate(truncate) return fh def __enter__(self): return self.acquire() def __exit__(self, type_, value, tb): self.release()