[PATCH 2 of 4 marmoute-reviewed] lock: turn prepinherit/reacquire into a single context manager

[Siddharth Agarwal]

# HG changeset patch
# User Siddharth Agarwal <sid0 at fb.com>
# Date 1444014170 25200
#      Sun Oct 04 20:02:50 2015 -0700
# Node ID 7ec3a1cc6f54cb2f94a199b2381d724ec071ba5d
# Parent  6c6d3c21524e6154db6c4f4816d345986f24dc9b
lock: turn prepinherit/reacquire into a single context manager

Review feedback from Pierre-Yves David. This makes the overall code cleaner and
less error-prone, and makes a previously explicitly checked error state
impossible.

diff --git a/mercurial/lock.py b/mercurial/lock.py
--- a/mercurial/lock.py
+++ b/mercurial/lock.py
@@ -7,6 +7,7 @@
 
 from __future__ import absolute_import
 
+import contextlib
 import errno
 import os
 import socket
@@ -171,19 +172,20 @@ class lock(object):
         locker = self._readlock()
         return self._testlock(locker)
 
-    def prepinherit(self):
-        """prepare for the lock to be inherited by a Mercurial subprocess
+    @contextlib.contextmanager
+    def inherit(self):
+        """context for the lock to be inherited by a Mercurial subprocess.
 
-        Returns a string that will be recognized by the lock in the
-        subprocess. Communicating this string to the subprocess needs to be done
-        separately -- typically by an environment variable.
+        Yields a string that will be recognized by the lock in the subprocess.
+        Communicating this string to the subprocess needs to be done separately
+        -- typically by an environment variable.
         """
         if not self.held:
             raise error.LockInheritanceContractViolation(
-                'prepinherit can only be called while lock is held')
+                'inherit can only be called while lock is held')
         if self._inherited:
             raise error.LockInheritanceContractViolation(
-                'prepinherit cannot be called while lock is already inherited')
+                'inherit cannot be called while lock is already inherited')
         if self.releasefn:
             self.releasefn()
         if self._parentheld:
@@ -191,15 +193,12 @@ class lock(object):
         else:
             lockname = '%s:%s' % (lock._host, self.pid)
         self._inherited = True
-        return lockname
-
-    def reacquire(self):
-        if not self._inherited:
-            raise error.LockInheritanceContractViolation(
-                'reacquire can only be called after prepinherit')
-        if self.acquirefn:
-            self.acquirefn()
-        self._inherited = False
+        try:
+            yield lockname
+        finally:
+            if self.acquirefn:
+                self.acquirefn()
+            self._inherited = False
 
     def release(self):
         """release the lock and execute callback function if any
diff --git a/tests/test-lock.py b/tests/test-lock.py
--- a/tests/test-lock.py
+++ b/tests/test-lock.py
@@ -158,23 +158,22 @@ class testlock(unittest.TestCase):
         parentstate.assertacquirecalled(True)
 
         # set up lock inheritance
-        lockname = parentlock.prepinherit()
-        parentstate.assertreleasecalled(True)
-        parentstate.assertpostreleasecalled(False)
-        parentstate.assertlockexists(True)
+        with parentlock.inherit() as lockname:
+            parentstate.assertreleasecalled(True)
+            parentstate.assertpostreleasecalled(False)
+            parentstate.assertlockexists(True)
 
-        childstate = teststate(self, d, pidoffset=1)
-        childlock = childstate.makelock(parentlock=lockname)
-        childstate.assertacquirecalled(True)
+            childstate = teststate(self, d, pidoffset=1)
+            childlock = childstate.makelock(parentlock=lockname)
+            childstate.assertacquirecalled(True)
 
-        # release the child lock -- the lock file should still exist on disk
-        childlock.release()
-        childstate.assertreleasecalled(True)
-        childstate.assertpostreleasecalled(True)
-        childstate.assertlockexists(True)
+            childlock.release()
+            childstate.assertreleasecalled(True)
+            childstate.assertpostreleasecalled(True)
+            childstate.assertlockexists(True)
 
-        parentstate.resetacquirefn()
-        parentlock.reacquire()
+            parentstate.resetacquirefn()
+
         parentstate.assertacquirecalled(True)
 
         parentlock.release()
@@ -188,39 +187,39 @@ class testlock(unittest.TestCase):
         lock0 = state0.makelock()
         state0.assertacquirecalled(True)
 
-        lock0name = lock0.prepinherit()
-        state0.assertreleasecalled(True)
-        state0.assertpostreleasecalled(False)
-        state0.assertlockexists(True)
+        with lock0.inherit() as lock0name:
+            state0.assertreleasecalled(True)
+            state0.assertpostreleasecalled(False)
+            state0.assertlockexists(True)
 
-        state1 = teststate(self, d, pidoffset=1)
-        lock1 = state1.makelock(parentlock=lock0name)
-        state1.assertacquirecalled(True)
+            state1 = teststate(self, d, pidoffset=1)
+            lock1 = state1.makelock(parentlock=lock0name)
+            state1.assertacquirecalled(True)
 
-        # from within lock1, acquire another lock
-        lock1name = lock1.prepinherit()
-        # since the file on disk is lock0's this should have the same name
-        self.assertEqual(lock0name, lock1name)
+            # from within lock1, acquire another lock
+            with lock1.inherit() as lock1name:
+                # since the file on disk is lock0's this should have the same
+                # name
+                self.assertEqual(lock0name, lock1name)
 
-        state2 = teststate(self, d, pidoffset=2)
-        lock2 = state2.makelock(parentlock=lock1name)
-        state2.assertacquirecalled(True)
+                state2 = teststate(self, d, pidoffset=2)
+                lock2 = state2.makelock(parentlock=lock1name)
+                state2.assertacquirecalled(True)
 
-        lock2.release()
-        state2.assertreleasecalled(True)
-        state2.assertpostreleasecalled(True)
-        state2.assertlockexists(True)
+                lock2.release()
+                state2.assertreleasecalled(True)
+                state2.assertpostreleasecalled(True)
+                state2.assertlockexists(True)
 
-        state1.resetacquirefn()
-        lock1.reacquire()
-        state1.assertacquirecalled(True)
+                state1.resetacquirefn()
 
-        lock1.release()
-        state1.assertreleasecalled(True)
-        state1.assertpostreleasecalled(True)
-        state1.assertlockexists(True)
+            state1.assertacquirecalled(True)
 
-        lock0.reacquire()
+            lock1.release()
+            state1.assertreleasecalled(True)
+            state1.assertpostreleasecalled(True)
+            state1.assertlockexists(True)
+
         lock0.release()
 
     def testinheritlockfork(self):
@@ -230,26 +229,25 @@ class testlock(unittest.TestCase):
         parentstate.assertacquirecalled(True)
 
         # set up lock inheritance
-        lockname = parentlock.prepinherit()
-        childstate = teststate(self, d, pidoffset=1)
-        childlock = childstate.makelock(parentlock=lockname)
-        childstate.assertacquirecalled(True)
+        with parentlock.inherit() as lockname:
+            childstate = teststate(self, d, pidoffset=1)
+            childlock = childstate.makelock(parentlock=lockname)
+            childstate.assertacquirecalled(True)
 
-        # fork the child lock
-        forkchildlock = copy.deepcopy(childlock)
-        forkchildlock._pidoffset += 1
-        forkchildlock.release()
-        childstate.assertreleasecalled(False)
-        childstate.assertpostreleasecalled(False)
-        childstate.assertlockexists(True)
+            # fork the child lock
+            forkchildlock = copy.deepcopy(childlock)
+            forkchildlock._pidoffset += 1
+            forkchildlock.release()
+            childstate.assertreleasecalled(False)
+            childstate.assertpostreleasecalled(False)
+            childstate.assertlockexists(True)
 
-        # release the child lock
-        childlock.release()
-        childstate.assertreleasecalled(True)
-        childstate.assertpostreleasecalled(True)
-        childstate.assertlockexists(True)
+            # release the child lock
+            childlock.release()
+            childstate.assertreleasecalled(True)
+            childstate.assertpostreleasecalled(True)
+            childstate.assertlockexists(True)
 
-        parentlock.reacquire()
         parentlock.release()
 
 if __name__ == '__main__':