Repository: cocotb/cocotb

Issue:
Add support for using Lock objects with async with
It would be great to add support for spelling
```python
await lock.acquire()
try:
    # do some stuff
finally:
    lock.release()
```
as
```python
async with lock:
    # do some stuff
```
This would be an easy first issue for someone contributing to cocotb, and would involve implementing `__aenter__` and `__exit__`, and writing a simple test.


Discussion:

Resolving patch:
diff --git a/cocotb/triggers.py b/cocotb/triggers.py
--- a/cocotb/triggers.py
+++ b/cocotb/triggers.py
@@ -456,13 +456,21 @@ def __repr__(self):
 class Lock(object):
     """Lock primitive (not re-entrant).
 
-    This should be used as::
+    This can be used as::
 
-        yield lock.acquire()
+        await lock.acquire()
         try:
             # do some stuff
         finally:
             lock.release()
+
+    .. versionchanged:: 1.4
+
+        The lock can be used as an asynchronous context manager in an
+        :keyword:`async with` statement::
+
+            async with lock:
+                # do some stuff
     """
 
     def __init__(self, name=None):
@@ -518,6 +526,12 @@ def __nonzero__(self):
 
     __bool__ = __nonzero__
 
+    async def __aenter__(self):
+        return await self.acquire()
+
+    async def __aexit__(self, exc_type, exc, tb):
+        self.release()
+
 
 class NullTrigger(Trigger):
     """Fires immediately.


Test patch:
diff --git a/tests/test_cases/test_cocotb/test_cocotb.py b/tests/test_cases/test_cocotb/test_cocotb.py
--- a/tests/test_cases/test_cocotb/test_cocotb.py
+++ b/tests/test_cases/test_cocotb/test_cocotb.py
@@ -47,7 +47,7 @@
 import cocotb
 from cocotb.triggers import (Timer, Join, RisingEdge, FallingEdge, Edge,
                              ReadOnly, ReadWrite, ClockCycles, NextTimeStep,
-                             NullTrigger, Combine, Event, First, Trigger)
+                             NullTrigger, Combine, Event, First, Trigger, Lock)
 from cocotb.clock import Clock
 from cocotb.result import (
     TestFailure, TestError, TestSuccess, raise_error, create_error
@@ -1466,3 +1466,48 @@ async def test_ordering_1(dut):
     global last_ordered_test
     val, last_ordered_test = last_ordered_test, 1
     assert val == 2
+
+
+@cocotb.test()
+async def test_trigger_lock(dut):
+    """
+    Simple test that checks to see if context management is kept. The
+    resource value is checked at certain points if it equals the expected
+    amount, which is easily predictable if the context management is working.
+    """
+    resource = 0
+    lock = Lock()
+
+    async def co():
+        nonlocal resource
+        await Timer(10, "ns")
+        async with lock:
+            for i in range(4):
+                await Timer(10, "ns")
+                resource += 1
+
+    cocotb.fork(co())
+    async with lock:
+        for i in range(4):
+            resource += 1
+            await Timer(10, "ns")
+    assert resource == 4
+    await Timer(10, "ns")
+    async with lock:
+        assert resource == 8
+
+
+@cocotb.test(timeout_time=100, timeout_unit="ns")
+async def test_except_lock(dut):
+    """
+    Checks to see if exceptions cause the lock to be
+    released.
+    """
+    lock = Lock()
+    try:
+        async with lock:
+            raise RuntimeError()
+    except RuntimeError:
+        pass
+    async with lock:
+        pass