refactoring

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:
@@ -152,22 +150,26 @@ class BaseTaskPool:
             raise e
         return task_id
 
-    def _cancel_one(self, task_id: int, msg: str = None) -> None:
+    def _get_running_task(self, task_id: int) -> Task:
         try:
-            task = self._running[task_id]
+            return self._running[task_id]
         except KeyError:
             if self._cancelled.get(task_id):
                 raise exceptions.AlreadyCancelled(f"{self._task_name(task_id)} has already been cancelled")
             if self._ended.get(task_id):
                 raise exceptions.AlreadyFinished(f"{self._task_name(task_id)} has finished running")
             raise exceptions.InvalidTaskID(f"No task with ID {task_id} found in {self}")
-        task.cancel(msg=msg)
+
+    def _cancel_task(self, task_id: int, msg: str = None) -> None:
+        self._get_running_task(task_id).cancel(msg=msg)
 
     def cancel(self, *task_ids: int, msg: str = None) -> None:
-        for task_id in task_ids:
-            self._cancel_one(task_id, msg=msg)
+        tasks = [self._get_running_task(task_id) for task_id in task_ids]
+        for task in tasks:
+            task.cancel(msg=msg)
 
-    def cancel_all(self, msg: str = None) -> None:
+    async def cancel_all(self, msg: str = None) -> None:
+        await self._all_tasks_known_flag.wait()
         for task in self._running.values():
             task.cancel(msg=msg)