From ba0d5fca858e85ee323cbdbeecc3ae8375cd011e Mon Sep 17 00:00:00 2001 From: Daniil Fajnberg Date: Sun, 6 Feb 2022 13:49:57 +0100 Subject: [PATCH] refactoring/renaming --- src/asyncio_taskpool/pool.py | 20 +++++++++----------- 1 file changed, 9 insertions(+), 11 deletions(-) diff --git a/src/asyncio_taskpool/pool.py b/src/asyncio_taskpool/pool.py index 1344388..c7f4993 100644 --- a/src/asyncio_taskpool/pool.py +++ b/src/asyncio_taskpool/pool.py @@ -106,20 +106,18 @@ class BaseTaskPool: """Returns a standardized name for a task with a specific `task_id`.""" return f'{self}_Task-{task_id}' - async def _cancel_task(self, task_id: int, custom_callback: CancelCallbackT = None) -> None: + async def _task_cancellation(self, task_id: int, custom_callback: CancelCallbackT = None) -> None: log.debug("Cancelling %s ...", self._task_name(task_id)) - task = self._running.pop(task_id) - assert task is not None - self._cancelled[task_id] = task + self._cancelled[task_id] = self._running.pop(task_id) self._num_cancelled += 1 log.debug("Cancelled %s", self._task_name(task_id)) await _execute_function(custom_callback, args=(task_id, )) - async def _end_task(self, task_id: int, custom_callback: EndCallbackT = None) -> None: - task = self._running.pop(task_id, None) - if task is None: - task = self._cancelled.pop(task_id) - self._ended[task_id] = task + async def _task_ending(self, task_id: int, custom_callback: EndCallbackT = None) -> None: + try: + self._ended[task_id] = self._running.pop(task_id) + except KeyError: + self._ended[task_id] = self._cancelled.pop(task_id) self._num_ended += 1 self._enough_room.release() log.info("Ended %s", self._task_name(task_id)) @@ -131,9 +129,9 @@ class BaseTaskPool: try: return await awaitable except CancelledError: - await self._cancel_task(task_id, custom_callback=cancel_callback) + await self._task_cancellation(task_id, custom_callback=cancel_callback) finally: - await self._end_task(task_id, custom_callback=end_callback) + await self._task_ending(task_id, custom_callback=end_callback) async def _start_task(self, awaitable: Awaitable, ignore_closed: bool = False, end_callback: EndCallbackT = None, cancel_callback: CancelCallbackT = None) -> int: