codemp-sublime/src/TaskManager.py

from typing import Optional
import asyncio
from ..ext import sublime_asyncio as rt


class TaskManager:
    def __init__(self):
        self.tasks = []
        self.runtime = rt
        self.exit_handler_id = None

    def acquire(self, exit_handler):
        if self.exit_handler_id is None:
            # don't allow multiple exit handlers
            self.exit_handler_id = self.runtime.acquire(exit_handler)

        return self.exit_handler_id

    def release(self, at_exit):
        self.runtime.release(at_exit=at_exit, exit_handler_id=self.exit_handler_id)

    def dispatch(self, coro, name=None):
        self.runtime.dispatch(coro, self.store_named_lambda(name))

    def sync(self, coro):
        self.runtime.sync(coro)

    def remove_stopped(self):
        self.tasks = list(filter(lambda T: not T.cancelled(), self.tasks))

    def store(self, task, name=None):
        if name is not None:
            task.set_name(name)
        self.tasks.append(task)
        self.remove_stopped()

    def store_named_lambda(self, name=None):
        def _store(task):
            self.store(task, name)

        return _store

    def get_task(self, name) -> Optional[asyncio.Task]:
        return next((t for t in self.tasks if t.get_name() == name), None)

    def get_task_idx(self, name) -> Optional[int]:
        return next(
            (i for (i, t) in enumerate(self.tasks) if t.get_name() == name), None
        )

    def pop_task(self, name) -> Optional[asyncio.Task]:
        idx = self.get_task_idx(name)
        if id is not None:
            return self.tasks.pop(idx)
        return None

    async def _stop(self, task):
        task.cancel()  # cancelling a task, merely requests a cancellation.
        try:
            await task
        except asyncio.CancelledError:
            return

    def stop(self, name):
        t = self.get_task(name)
        if t is not None:
            self.runtime.dispatch(self._stop(t))

    def stop_all(self):
        for task in self.tasks:
            self.runtime.dispatch(self._stop(task))


# singleton instance
tm = TaskManager()