|
16 | 16 | 'isfile', |
17 | 17 | 'isdir', |
18 | 18 | 'find_duplicates', |
19 | | - 'print_fs'] |
| 19 | + 'print_fs', |
| 20 | + 'open_atomic_write'] |
20 | 21 |
|
| 22 | +import os |
21 | 23 | import sys |
22 | 24 | import stat |
23 | 25 | import six |
@@ -628,6 +630,46 @@ def print_dir(fs, path, levels=[]): |
628 | 630 | return dircount[0], filecount[0] |
629 | 631 |
|
630 | 632 |
|
| 633 | +class AtomicWriter(object): |
| 634 | + """Context manager to perform atomic writes""" |
| 635 | + |
| 636 | + def __init__(self, fs, path, mode='w'): |
| 637 | + self.fs = fs |
| 638 | + self.path = path |
| 639 | + self.mode = mode |
| 640 | + self.tmp_path = path + '~' |
| 641 | + self._f = None |
| 642 | + |
| 643 | + def __enter__(self): |
| 644 | + self._f = self.fs.open(self.tmp_path, self.mode) |
| 645 | + return self._f |
| 646 | + |
| 647 | + def __exit__(self, exc_type, exc_value, traceback): |
| 648 | + if exc_type is None: |
| 649 | + if self._f is not None: |
| 650 | + if hasattr('_f', 'flush'): |
| 651 | + self._f.flush() |
| 652 | + if hasattr(self._f, 'fileno'): |
| 653 | + os.fsync(self._f.fileno()) |
| 654 | + self._f.close() |
| 655 | + self._f = None |
| 656 | + self.fs.rename(self.tmp_path, self.path) |
| 657 | + else: |
| 658 | + if self._f is not None: |
| 659 | + self._f.close() |
| 660 | + |
| 661 | + |
| 662 | +def open_atomic_write(fs, path, mode='w'): |
| 663 | + """Open a file for 'atomic' writing |
| 664 | +
|
| 665 | + This returns a context manager which ensures that a file is written in its entirety or not at all. |
| 666 | +
|
| 667 | + """ |
| 668 | + return AtomicWriter(fs, path, mode=mode) |
| 669 | + |
| 670 | + |
| 671 | + |
| 672 | + |
631 | 673 | if __name__ == "__main__": |
632 | 674 | from fs.tempfs import TempFS |
633 | 675 | from six import b |
|
0 commit comments