Merge branch 'master' of https://github.com/labthings/python-labthings

Author: Joel Collins

src/labthings/core/lock.py

@@ -1,6 +1,7 @@
 from gevent.hub import getcurrent
 from gevent.lock import RLock as _RLock
 
+from contextlib import contextmanager
 import logging
 
 from .exceptions import LockError
@@ -31,6 +32,13 @@ def __init__(self, timeout=None, name=None):
         self.timeout = timeout
         self.name = name
 
+    @contextmanager
+    def __call__(self, timeout=sentinel, blocking=True):
+        result = self.acquire(timeout=timeout, blocking=blocking)
+        yield result
+        if result:
+            self.release()
+        
     def locked(self):
         return self._lock.locked()
 
@@ -82,6 +90,13 @@ def __init__(self, locks, timeout=None):
         self.locks = locks
         self.timeout = timeout
 
+    @contextmanager
+    def __call__(self, timeout=sentinel, blocking=True):
+        result = self.acquire(timeout=timeout, blocking=blocking)
+        yield result
+        if result:
+            self.release()
+        
     def acquire(self, blocking=True, timeout=sentinel):
         if timeout is sentinel:
             timeout = self.timeout

tests/test_core_lock.py

@@ -88,3 +88,26 @@ def test_rlock_block(this_lock):
 
     # Release lock, assert no owner
     assert not this_lock._is_owned()
+    
+def test_rlock_acquire_timeout(this_lock):
+    from labthings.core.exceptions import LockError
+
+    # Acquire lock
+    assert this_lock.acquire()
+
+    # Override owner to force acquisition failure
+    this_lock._owner = None
+    print(this_lock._owner)
+    # Assert not owner
+    assert not this_lock._is_owned()
+
+    # Assert acquisition fails using context manager
+    with pytest.raises(LockError):
+        with this_lock(timeout=0.01):
+            pass
+
+    # Force ownership
+    this_lock._owner = getcurrent()
+
+    # Release lock
+    this_lock.release()