diff --git a/src/cutecoin/tools/decorators.py b/src/cutecoin/tools/decorators.py index 3842dbd419aa04b4275fe5f303ac48a5e2d8b443..169bdae169dc49e22b1a64a39dcd0cbcb7a10cf6 100644 --- a/src/cutecoin/tools/decorators.py +++ b/src/cutecoin/tools/decorators.py @@ -13,15 +13,24 @@ def cancel_once_task(object, fn): def once_at_a_time(fn): @functools.wraps(fn) def wrapper(*args, **kwargs): + def task_done(task): + try: + args[0].__tasks.pop(fn.__name__) + except KeyError: + logging.debug("Task already removed") + if getattr(args[0], "__tasks", None) is None: setattr(args[0], "__tasks", {}) if fn.__name__ in args[0].__tasks: if not args[0].__tasks[fn.__name__].done(): args[0].__tasks[fn.__name__].cancel() + try: args[0].__tasks[fn.__name__] = fn(*args, **kwargs) + args[0].__tasks[fn.__name__].add_done_callback(task_done) except asyncio.CancelledError: logging.debug("Cancelled asyncified : {0}".format(fn.__name__)) + return args[0].__tasks[fn.__name__] return wrapper