codekingpro/portable-devtools

1import asyncio2import os3import sys4import time5from collections.abc import Coroutine6from collections.abc import Iterator7from contextlib import contextmanager8 9from mitmproxy.utils import human10 11_KEEP_ALIVE = set()12 13 14def create_task(15    coro: Coroutine,16    *,17    name: str,18    keep_ref: bool,19    client: tuple | None = None,20) -> asyncio.Task:21    """22    Wrapper around `asyncio.create_task`.23 24    - Use `keep_ref` to keep an internal reference.25      This ensures that the task is not garbage collected mid-execution if no other reference is kept.26    - Use `client` to pass the client address as additional debug info on the task.27    """28    t = asyncio.create_task(coro)  # noqa: TID25129    set_task_debug_info(t, name=name, client=client)30    if keep_ref and not t.done():31        # The event loop only keeps weak references to tasks.32        # A task that isn’t referenced elsewhere may get garbage collected at any time, even before it’s done.33        _KEEP_ALIVE.add(t)34        t.add_done_callback(_KEEP_ALIVE.discard)35    return t36 37 38def set_task_debug_info(39    task: asyncio.Task,40    *,41    name: str,42    client: tuple | None = None,43) -> None:44    """Set debug info for an externally-spawned task."""45    task.created = time.time()  # type: ignore46    if __debug__ is True and (test := os.environ.get("PYTEST_CURRENT_TEST", None)):47        name = f"{name} [created in {test}]"48    task.set_name(name)49    if client:50        task.client = client  # type: ignore51 52 53def set_current_task_debug_info(54    *,55    name: str,56    client: tuple | None = None,57) -> None:58    """Set debug info for the current task."""59    task = asyncio.current_task()60    assert task61    set_task_debug_info(task, name=name, client=client)62 63 64def task_repr(task: asyncio.Task) -> str:65    """Get a task representation with debug info."""66    name = task.get_name()67    a: float = getattr(task, "created", 0)68    if a:69        age = f" (age: {time.time() - a:.0f}s)"70    else:71        age = ""72    client = getattr(task, "client", "")73    if client:74        client = f"{human.format_address(client)}: "75    return f"{client}{name}{age}"76 77 78@contextmanager79def install_exception_handler(handler) -> Iterator[None]:80    loop = asyncio.get_running_loop()81    existing = loop.get_exception_handler()82    loop.set_exception_handler(handler)83    try:84        yield85    finally:86        loop.set_exception_handler(existing)87 88 89@contextmanager90def set_eager_task_factory() -> Iterator[None]:91    loop = asyncio.get_running_loop()92    if sys.version_info < (3, 12):  # pragma: no cover93        yield94    else:95        existing = loop.get_task_factory()96        loop.set_task_factory(asyncio.eager_task_factory)  # type: ignore97        try:98            yield99        finally:100            loop.set_task_factory(existing)101