""" Various object wrappers and abstraction layers Copyright (C) 2020 nocturn9x Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. """ from giambio.traps import join, cancel, event_set, event_wait from heapq import heappop, heappush, heapify from giambio.exceptions import GiambioError from dataclasses import dataclass, field import typing @dataclass class Task: """ A simple wrapper around a coroutine object """ coroutine: typing.Coroutine name: str pool: "giambio.context.TaskManager" cancelled: bool = False exc: BaseException = None result: object = None finished: bool = False status: str = "init" steps: int = 0 last_io: tuple = () joiners: list = field(default_factory=list) joined: bool = False cancel_pending: bool = False sleep_start: float = 0.0 next_deadline: float = 0.0 def run(self, what: object = None): """ Simple abstraction layer over coroutines' ``send`` method """ return self.coroutine.send(what) def throw(self, err: Exception): """ Simple abstraction layer over coroutines ``throw`` method """ return self.coroutine.throw(err) async def join(self): """ Joins the task """ res = await join(self) if self.exc: raise self.exc return res async def cancel(self): """ Cancels the task """ await cancel(self) def __hash__(self): return hash(self.coroutine) def done(self): return self.exc or self.finished or self.cancelled class Event: """ A class designed similarly to threading.Event """ def __init__(self): """ Object constructor """ self.set = False self.waiters = [] self.event_caught = False async def trigger(self): """ Sets the event, waking up all tasks that called pause() on us """ if self.set: # This is set by the event loop internally raise GiambioError("The event has already been set") await event_set(self) async def wait(self): """ Waits until the event is set """ await event_wait(self) class TimeQueue: """ An abstraction layer over a heap queue based on time. This is where sleeping tasks will be put when they are not running """ def __init__(self, clock): """ Object constructor """ self.clock = clock self.sequence = 0 # The sequence number handles the race condition # of two tasks with identical deadlines acting # as a tie breaker self.container = [] def __contains__(self, item): for i in self.container: if i[2] == item: return True return False def discard(self, item): for i in self.container: if i[2] == item: self.container.remove(i) heapify(self.container) return def __iter__(self): return self def __next__(self): try: return self.get() except IndexError: raise StopIteration from None def __getitem__(self, item): return self.container.__getitem__(item) def __bool__(self): return bool(self.container) def __repr__(self): return f"TimeQueue({self.container}, clock={self.clock})" def put(self, item, amount: float): """ Pushes an item onto the queue with its unique time amount and ID """ heappush(self.container, (self.clock() + amount, self.sequence, item)) self.sequence += 1 def get(self): """ Gets the first task that is meant to run """ return heappop(self.container)[2] class DeadlinesQueue(TimeQueue): """ An ordered queue for storing tasks deadlines """ def __init__(self): """ Object constructor """ super().__init__(None) self.pools = set() def __contains__(self, item): return super().__contains__(item) def __iter__(self): return super().__iter__() def __next__(self): return super().__next__() def __getitem__(self, item): return super().__getitem__(item) def __bool__(self): return super().__bool__() def __repr__(self): return f"DeadlinesQueue({self.container})" def put(self, amount: float, pool): """ Pushes a deadline (timeout) onto the queue with its associated pool """ if pool not in self.pools: self.pools.add(pool) heappush(self.container, (amount, self.sequence, pool)) def get(self): """ Gets the first task that is meant to run """ d = heappop(self.container) self.pools.discard(d[2]) return d