From 8bc2aba1d41b91dbfc699950e80ed9e6d748a051 Mon Sep 17 00:00:00 2001 From: BobTheBuidler <70677534+BobTheBuidler@users.noreply.github.com> Date: Sat, 23 Nov 2024 04:22:28 -0400 Subject: [PATCH] feat: speed up SmartTask and SmartFuture init (#433) * feat: speed up SmartTask and SmartFuture init * chore: `black .` --------- Co-authored-by: github-actions[bot] --- a_sync/_smart.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/a_sync/_smart.py b/a_sync/_smart.py index b6e901d7..e61a215e 100644 --- a/a_sync/_smart.py +++ b/a_sync/_smart.py @@ -6,6 +6,7 @@ """ import asyncio +import contextvars import logging import warnings import weakref @@ -197,7 +198,9 @@ def __init__( if key: self._key = key self._waiters = weakref.WeakSet() - self.add_done_callback(SmartFuture._self_done_cleanup_callback) + self._callbacks.append( + (SmartFuture._self_done_cleanup_callback, contextvars.copy_context()) + ) def __repr__(self): return f"<{type(self).__name__} key={self._key} waiters={self.num_waiters} {self._state}>" @@ -298,7 +301,9 @@ def __init__( """ asyncio.Task.__init__(self, coro, loop=loop, name=name) self._waiters: Set["asyncio.Task[T]"] = set() - self.add_done_callback(SmartTask._self_done_cleanup_callback) + self._callbacks.append( + (SmartTask._self_done_cleanup_callback, contextvars.copy_context()) + ) def smart_task_factory(