codekingpro/portable-devtools

1from __future__ import annotations2 3import asyncio4import concurrent.futures5import contextvars6import inspect7import sys8import types9from collections.abc import Awaitable, Coroutine, Generator10from typing import TypeVar, cast11 12T = TypeVar("T")13AnyFuture = asyncio.Future | concurrent.futures.Future14 15CONTEXT_NOT_SUPPORTED = sys.version_info < (3, 11)16EAGER_NOT_SUPPORTED = sys.version_info < (3, 12)17 18 19def _get_loop(fut: asyncio.Future) -> asyncio.AbstractEventLoop:20    # Tries to call Future.get_loop() if it's available.21    # Otherwise fallbacks to using the old '_loop' property.22    try:23        get_loop = fut.get_loop24    except AttributeError:25        pass26    else:27        return get_loop()28    return fut._loop29 30 31def _convert_future_exc(exc: BaseException) -> BaseException:32    exc_class = type(exc)33    if exc_class is concurrent.futures.CancelledError:34        return asyncio.CancelledError(*exc.args)35    elif exc_class is concurrent.futures.TimeoutError:36        return asyncio.TimeoutError(*exc.args)37    elif exc_class is concurrent.futures.InvalidStateError:38        return asyncio.InvalidStateError(*exc.args)39    else:40        return exc41 42 43def _set_concurrent_future_state(44    concurrent: concurrent.futures.Future,45    source: AnyFuture,46) -> None:47    """Copy state from a future to a concurrent.futures.Future."""48    assert source.done()49    if source.cancelled():50        concurrent.cancel()51    if not concurrent.set_running_or_notify_cancel():52        return53    exception = source.exception()54    if exception is not None:55        concurrent.set_exception(_convert_future_exc(exception))56    else:57        result = source.result()58        concurrent.set_result(result)59 60 61def _copy_future_state(source: AnyFuture, dest: asyncio.Future) -> None:62    """Internal helper to copy state from another Future.63 64    The other Future may be a concurrent.futures.Future.65    """66    if dest.done():67        return68    assert source.done()69    if dest.cancelled():70        return71    if source.cancelled():72        dest.cancel()73    else:74        exception = source.exception()75        if exception is not None:76            dest.set_exception(_convert_future_exc(exception))77        else:78            result = source.result()79            dest.set_result(result)80 81 82def _chain_future(source: AnyFuture, destination: AnyFuture) -> None:83    """Chain two futures so that when one completes, so does the other.84 85    The result (or exception) of source will be copied to destination.86    If destination is cancelled, source gets cancelled too.87    Compatible with both asyncio.Future and concurrent.futures.Future.88    """89    if not asyncio.isfuture(source) and not isinstance(90        source, concurrent.futures.Future91    ):92        raise TypeError("A future is required for source argument")93    if not asyncio.isfuture(destination) and not isinstance(94        destination, concurrent.futures.Future95    ):96        raise TypeError("A future is required for destination argument")97    source_loop = _get_loop(source) if asyncio.isfuture(source) else None98    dest_loop = _get_loop(destination) if asyncio.isfuture(destination) else None99 100    def _set_state(future: AnyFuture, other: AnyFuture) -> None:101        if asyncio.isfuture(future):102            _copy_future_state(other, future)103        else:104            _set_concurrent_future_state(future, other)105 106    def _call_check_cancel(destination: AnyFuture) -> None:107        if destination.cancelled():108            if source_loop is None or source_loop is dest_loop:109                source.cancel()110            else:111                source_loop.call_soon_threadsafe(source.cancel)112 113    def _call_set_state(source: AnyFuture) -> None:114        if destination.cancelled() and dest_loop is not None and dest_loop.is_closed():115            return116        if dest_loop is None or dest_loop is source_loop:117            _set_state(destination, source)118        else:119            if dest_loop.is_closed():120                return121            dest_loop.call_soon_threadsafe(_set_state, destination, source)122 123    destination.add_done_callback(_call_check_cancel)124    source.add_done_callback(_call_set_state)125 126 127def chain_future(source: AnyFuture, destination: AnyFuture) -> AnyFuture:128    # adapted from asyncio.run_coroutine_threadsafe129    try:130        _chain_future(source, destination)131        return destination132    except (SystemExit, KeyboardInterrupt):133        raise134    except BaseException as exc:135        if isinstance(destination, concurrent.futures.Future):136            if destination.set_running_or_notify_cancel():137                destination.set_exception(exc)138        else:139            destination.set_exception(exc)140        raise141 142 143def _ensure_future(144    coro_or_future: Coroutine[None, None, T] | Awaitable[T],145    *,146    loop: asyncio.AbstractEventLoop,147    name: str | None = None,148    context: contextvars.Context | None = None,149    lazy: bool = True,150) -> asyncio.Task[T]:151    called_wrap_awaitable = False152    if not asyncio.iscoroutine(coro_or_future):153        if inspect.isawaitable(coro_or_future):154            coro_or_future = cast(155                Coroutine[None, None, T], _wrap_awaitable(coro_or_future)156            )157            called_wrap_awaitable = True158        else:159            raise TypeError(160                "An asyncio.Future, a coroutine or an awaitable is required."161                f" Got {type(coro_or_future).__name__} instead."162            )163 164    try:165        if CONTEXT_NOT_SUPPORTED:166            return loop.create_task(coro_or_future, name=name)167        elif EAGER_NOT_SUPPORTED or lazy:168            return loop.create_task(coro_or_future, name=name, context=context)169        else:170            return asyncio.eager_task_factory(171                loop, coro_or_future, name=name, context=context172            )173    except RuntimeError:174        if not called_wrap_awaitable:175            coro_or_future.close()176        raise177 178 179@types.coroutine180def _wrap_awaitable(awaitable: Awaitable[T]) -> Generator[None, None, T]:181    """Helper for asyncio.ensure_future().182 183    Wraps awaitable (an object with __await__) into a coroutine184    that will later be wrapped in a Task by ensure_future().185    """186    return (yield from awaitable.__await__())187 188 189def run_coroutine_threadsafe(190    coro: Coroutine[None, None, T],191    loop: asyncio.AbstractEventLoop,192    *,193    lazy: bool,194    name: str | None = None,195    context: contextvars.Context | None = None,196) -> asyncio.Future[T]:197    """Submit a coroutine object to a given event loop.198 199    Return an asyncio.Future to access the result.200    """201 202    if asyncio._get_running_loop() is loop:203        return _ensure_future(coro, loop=loop, name=name, context=context, lazy=lazy)204    else:205        future: asyncio.Future[T] = asyncio.Future(loop=loop)206 207        def callback() -> None:208            try:209                chain_future(210                    _ensure_future(coro, loop=loop, name=name, context=context),211                    future,212                )213            except (SystemExit, KeyboardInterrupt):214                raise215            except BaseException as exc:216                future.set_exception(exc)217                raise218 219        loop.call_soon_threadsafe(callback, context=context)220        return future221