aiosched/aiosched/context.py

"""
aiosched: Yet another Python async scheduler

Copyright (C) 2022 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

   https: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 aiosched.task import Task
from aiosched.internals.syscalls import (
    spawn,
    wait,
    cancel,
    join,
    current_task
)
from typing import Any, Coroutine, Callable


class TaskContext:
    """
    An asynchronous context manager that automatically waits
    for all tasks spawned within it and cancels itself when
    an exception occurs. Contexts can be nested and will
    cancel inner ones if an exception is raised inside them
    """

    def __init__(self, silent: bool = False, gather: bool = True, timeout: int | float = 0.0) -> None:
        """
        Object constructor
        """

        # All the tasks that belong to this context
        self.tasks: list[Task] = []
        # Whether we have been cancelled or not
        self.cancelled: bool = False
        # The context's entry point (needed to disguise ourselves as a task ;))
        self.entry_point: Task | TaskContext | None = None
        # Do we ignore exceptions?
        self.silent: bool = silent
        # Do we gather multiple exceptions from
        # children tasks?
        self.gather: bool = gather   # TODO: Implement
        # For how long do we allow tasks inside us
        # to run?
        self.timeout: int | float = timeout   # TODO: Implement
        # Have we crashed?
        self.error: BaseException | None = None

    async def spawn(
            self, func: Callable[..., Coroutine[Any, Any, Any]], *args, **kwargs
    ) -> Task:
        """
        Spawns a child task
        """

        task = await spawn(func, *args, **kwargs)
        task.context = self
        self.tasks.append(task)
        await join(task)
        return task

    async def __aenter__(self):
        """
        Implements the asynchronous context manager interface
        """

        self.entry_point = await current_task()
        return self

    def __eq__(self, other):
        """
        Implements self == other
        """

        if isinstance(other, TaskContext):
            return super().__eq__(other)
        elif isinstance(other, Task):
            return other == self.entry_point
        return False

    async def __aexit__(self, exc_type: Exception, exc: Exception, tb):
        """
        Implements the asynchronous context manager interface, waiting
        for all the tasks spawned inside the context and handling
        exceptions
        """

        try:
            for task in self.tasks:
                # This forces the interpreter to stop at the
                # end of the block and wait for all
                # children to exit
                if task is self.entry_point:
                    # We don't wait on the entry
                    # point because that's us!
                    # Besides, even if we tried,
                    # wait() would raise an error
                    # to avoid a deadlock
                    continue
                await wait(task)
        except BaseException as exc:
            await self.cancel(False)
            self.error = exc
        finally:
            self.entry_point.propagate = True
            if self.silent:
                return
            if self.entry_point.exc:
                raise self.entry_point.exc

    # Task method wrappers

    async def cancel(self, propagate: bool = True):
        """
        Cancels the entire context, iterating over all
        of its tasks (which includes inner contexts)
        and cancelling them
        """

        for task in self.tasks:
            if task is self.entry_point:
                continue
            if isinstance(task, Task):
                await cancel(task)
            else:
                task: TaskContext
                await task.cancel(propagate)
        self.cancelled = True
        if propagate:
            if isinstance(self.entry_point, Task):
                await cancel(self.entry_point)
            else:
                self.entry_point: TaskContext
                await self.entry_point.cancel(propagate)

    def done(self) -> bool:
        """
        Returns whether all the tasks inside the
        context have exited
        """

        for task in self.tasks:
            if not task.done():
                return False
        return True

    def __hash__(self):
        return self.entry_point.__hash__()

    def run(self, what: Any | None = None):
        return self.entry_point.run(what)

    def __del__(self):
        """
        Context destructor
        """

        for task in self.tasks:
            task.__del__()

    def __repr__(self):
        """
        Implements repr(self)
        """

        result = "TaskContext(["
        for i, task in enumerate(self.tasks):
            result += repr(task)
            if i < len(self.tasks) - 1:
                result += ", "
        result += "])"
        return result