primitives/condition.py

   1 # condition.py
   2
   3 # Copyright (c) 2018-2020 Peter Hinch
   4 # Released under the MIT License (MIT) - see LICENSE file
   5
   6 try:
   7     import uasyncio as asyncio
   8 except ImportError:
   9     import asyncio
  10
  11 # Condition class
  12 # from primitives.condition import Condition
  13
  14 class Condition():
  15     def __init__(self, lock=None):
  16         self.lock = asyncio.Lock() if lock is None else lock
  17         self.events = []
  18
  19     async def acquire(self):
  20         await self.lock.acquire()
  21
  22 # enable this syntax:
  23 # with await condition [as cond]:
  24     def __await__(self):
  25         await self.lock.acquire()
  26         return self
  27
  28     __iter__ = __await__
  29
  30     def __enter__(self):
  31         return self
  32
  33     def __exit__(self, *_):
  34         self.lock.release()
  35
  36     def locked(self):
  37         return self.lock.locked()
  38
  39     def release(self):
  40         self.lock.release()  # Will raise RuntimeError if not locked
  41
  42     def notify(self, n=1):  # Caller controls lock
  43         if not self.lock.locked():
  44             raise RuntimeError('Condition notify with lock not acquired.')
  45         for _ in range(min(n, len(self.events))):
  46             ev = self.events.pop()
  47             ev.set()
  48
  49     def notify_all(self):
  50         self.notify(len(self.events))
  51
  52     async def wait(self):
  53         if not self.lock.locked():
  54             raise RuntimeError('Condition wait with lock not acquired.')
  55         ev = asyncio.Event()
  56         self.events.append(ev)
  57         self.lock.release()
  58         await ev.wait()
  59         await self.lock.acquire()
  60         assert ev not in self.events, 'condition wait assertion fail'
  61         return True  # CPython compatibility
  62
  63     async def wait_for(self, predicate):
  64         result = predicate()
  65         while not result:
  66             await self.wait()
  67             result = predicate()
  68         return result