From af6727951aa5f13845978d58559bae94d21325ce Mon Sep 17 00:00:00 2001
From: Richard Shaw <richard@phas.ubc.ca>
Date: Wed, 10 Apr 2019 16:44:15 -0700
Subject: [PATCH] feat: context manager for lockfiles

A context manager for wrapping around file creation operations that
automates the creation and cleaning of lock files, as well as removal of
the temporary file when there is a failure.
---
 caput/misc.py            | 73 +++++++++++++++++++++++++++++++
 caput/tests/test_misc.py | 93 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 166 insertions(+)
 create mode 100644 caput/tests/test_misc.py

diff --git a/caput/misc.py b/caput/misc.py
index 2b069487..89a4b728 100644
--- a/caput/misc.py
+++ b/caput/misc.py
@@ -16,6 +16,7 @@
 from future.builtins.disabled import *  # noqa  pylint: disable=W0401, W0614
 # === End Python 2/3 compatibility
 
+import os
 from past.builtins import basestring
 import numpy as np
 
@@ -118,3 +119,75 @@ def open_h5py_mpi(f, mode, comm=None):
     fh.is_mpi = (fh.file.driver == 'mpio')
 
     return fh
+
+
+class lock_file(object):
+    """Manage a lock file around a file creation operation.
+
+    Parameters
+    ----------
+    filename : str
+        Final name for the file.
+    preserve : bool, optional
+        Keep the temporary file in the event of failure.
+    comm : MPI.COMM, optional
+        If present only rank=0 will create/remove the lock file and move the
+        file.
+
+    Returns
+    -------
+    tmp_name : str
+        File name to use in the locked block.
+
+    Example
+    -------
+
+    >>> with lock_file('file_to_create.h5') as fname:
+    ...     container.save(fname)
+    ...
+    """
+
+    def __init__(self, name, preserve=False, comm=None):
+
+        from . import mpiutil
+
+        if comm is not None and not hasattr(comm, 'rank'):
+            raise ValueError('comm argument does not seem to be an MPI communicator.')
+
+        self.name = name
+        self.rank0 = mpiutil.rank0 if comm is None else comm.rank == 0
+        self.preserve = preserve
+
+    def __enter__(self):
+
+        if self.rank0:
+            with open(self.lockfile, 'w+') as fh:
+                fh.write("")
+
+        return self.tmpfile
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+
+        if self.rank0:
+
+            # Check if exception was raised and delete the temp file if needed
+            if exc_type is not None:
+                if not self.preserve:
+                    os.remove(self.tmpfile)
+            # Otherwise things were successful and we should move the file over
+            else:
+                os.rename(self.tmpfile, self.name)
+
+            # Finally remove the lock file
+            os.remove(self.lockfile)
+
+        return False
+
+    @property
+    def tmpfile(self):
+        base, fname = os.path.split(self.name)
+        return os.path.join(base, '.' + fname)
+
+    @property
+    def lockfile(self):
+        return self.tmpfile + '.lock'
diff --git a/caput/tests/test_misc.py b/caput/tests/test_misc.py
new file mode 100644
index 00000000..6f2f7f95
--- /dev/null
+++ b/caput/tests/test_misc.py
@@ -0,0 +1,93 @@
+"""Test the miscellaneous tools."""
+
+import unittest
+import tempfile
+import os
+import shutil
+
+from caput import misc
+
+class TestLock(unittest.TestCase):
+
+    def setUp(self):
+        self.dir = tempfile.mkdtemp()
+
+    def test_lock_new(self):
+        """Test the normal behaviour"""
+
+        base = 'newfile.dat'
+        newfile_name = os.path.join(self.dir, base)
+        lockfile_name = os.path.join(self.dir, '.' + base + '.lock')
+
+        with misc.lock_file(newfile_name) as fname:
+
+            # Check lock file has been created
+            self.assertTrue(os.path.exists(lockfile_name))
+
+            # Create a stub file
+            with open(fname, 'w+') as fh:
+                fh.write("hello")
+
+            # Check the file exists only at the temporary path
+            self.assertTrue(os.path.exists(fname))
+            self.assertFalse(os.path.exists(newfile_name))
+
+        # Check the file exists at the final path and the lock file removed
+        self.assertTrue(os.path.exists(newfile_name))
+        self.assertFalse(os.path.exists(lockfile_name))
+
+    def test_lock_exception(self):
+        """Check what happens in an exception"""
+
+        base = 'newfile2.dat'
+        newfile_name = os.path.join(self.dir, base)
+        lockfile_name = os.path.join(self.dir, '.' + base + '.lock')
+
+        try:
+            with misc.lock_file(newfile_name) as fname:
+
+                # Create a stub file
+                with open(fname, 'w+') as fh:
+                    fh.write("hello")
+
+                raise RuntimeError("Test error")
+
+        except:
+            pass
+
+        # Check that neither the file, nor its lock exists
+        self.assertFalse(os.path.exists(newfile_name))
+        self.assertFalse(os.path.exists(lockfile_name))
+
+    def test_lock_exception_preserve(self):
+        """Check what happens in an exception when asked to preserve the temp file"""
+
+        base = 'newfile3.dat'
+        newfile_name = os.path.join(self.dir, base)
+        lockfile_name = os.path.join(self.dir, '.' + base + '.lock')
+        tmpfile_name = os.path.join(self.dir, '.' + base)
+
+        try:
+            with misc.lock_file(newfile_name, preserve=True) as fname:
+
+                # Create a stub file
+                with open(fname, 'w+') as fh:
+                    fh.write("hello")
+
+                raise RuntimeError("Test error")
+
+        except:
+            pass
+
+        # Check that neither the file, nor its lock exists, but that the
+        # temporary file does
+        self.assertTrue(os.path.exists(tmpfile_name))
+        self.assertFalse(os.path.exists(newfile_name))
+        self.assertFalse(os.path.exists(lockfile_name))
+
+    def tearDown(self):
+        shutil.rmtree(self.dir)
+
+
+if __name__ == '__main__':
+    unittest.main()
\ No newline at end of file