diff --git a/atomicwrites/__init__.py b/atomicwrites/__init__.py index 0b9f92a..596ba1f 100644 --- a/atomicwrites/__init__.py +++ b/atomicwrites/__init__.py @@ -1,4 +1,5 @@ import contextlib +import errno import io import os import sys @@ -124,6 +125,12 @@ class AtomicWriter(object): :param overwrite: If set to false, an error is raised if ``path`` exists. Errors are only raised after the file has been written to. Either way, the operation is atomic. + :param path_generator: A generator function that takes the `path` as an + argument and returns the next path to attempt to create. If creation + fails, the generator will be repeatedly called to get the next path to + try writing until the write succeeds or the a previously attempted path + is generated again. A `ValueError` is raised if a previously attempted + path is provided by the generator. :param open_kwargs: Keyword-arguments to pass to the underlying :py:func:`open` call. This can be used to set the encoding when opening files in text-mode. @@ -133,7 +140,7 @@ class AtomicWriter(object): ''' def __init__(self, path, mode=DEFAULT_MODE, overwrite=False, - **open_kwargs): + path_generator=None, **open_kwargs): if 'a' in mode: raise ValueError( 'Appending to an existing file is not supported, because that ' @@ -153,7 +160,9 @@ def __init__(self, path, mode=DEFAULT_MODE, overwrite=False, self._path = path self._mode = mode self._overwrite = overwrite + self._path_generator = path_generator self._open_kwargs = open_kwargs + self.final_path = None def open(self): ''' @@ -169,10 +178,11 @@ def _open(self, get_fileobject): with get_fileobject(**self._open_kwargs) as f: yield f self.sync(f) - self.commit(f) + self.final_path = self.commit(f) success = True finally: if not success: + self.final_path = None try: self.rollback(f) except Exception: @@ -203,8 +213,31 @@ def commit(self, f): '''Move the temporary file to the target location.''' if self._overwrite: replace_atomic(f.name, self._path) + return self._path else: - move_atomic(f.name, self._path) + if self._path_generator is not None: + seen = set() + for path in self._path_generator(self._path): + if path in seen: + # avoid infinite loop if the path generator returns a + # path that was already attempted + raise ValueError( + 'path_generator must return unique values, but' + '{} was returned multiple times.'.format(path) + ) + seen.add(path) + try: + move_atomic(f.name, path) + except OSError as exc: + if exc.errno == errno.EEXIST: + pass + else: + raise + else: + return path + else: + move_atomic(f.name, self._path) + return self._path def rollback(self, f): '''Clean up all temporary resources.'''