generated from daniil-berg/boilerplate-py
huge update:
introduced meta tasks which are used by `_map()`; introduced task groups; ending with `gather_and_close()` now; pool unittests rewritten accordingly; two new helper classes
This commit is contained in:
75
src/asyncio_taskpool/group_register.py
Normal file
75
src/asyncio_taskpool/group_register.py
Normal file
@ -0,0 +1,75 @@
|
||||
__author__ = "Daniil Fajnberg"
|
||||
__copyright__ = "Copyright © 2022 Daniil Fajnberg"
|
||||
__license__ = """GNU LGPLv3.0
|
||||
|
||||
This file is part of asyncio-taskpool.
|
||||
|
||||
asyncio-taskpool is free software: you can redistribute it and/or modify it under the terms of
|
||||
version 3.0 of the GNU Lesser General Public License as published by the Free Software Foundation.
|
||||
|
||||
asyncio-taskpool is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY;
|
||||
without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
|
||||
See the GNU Lesser General Public License for more details.
|
||||
|
||||
You should have received a copy of the GNU Lesser General Public License along with asyncio-taskpool.
|
||||
If not, see <https://www.gnu.org/licenses/>."""
|
||||
|
||||
__doc__ = """
|
||||
This module contains the definition of the `TaskGroupRegister` class.
|
||||
"""
|
||||
|
||||
|
||||
from asyncio.locks import Lock
|
||||
from collections.abc import MutableSet
|
||||
from typing import Iterator, Set
|
||||
|
||||
|
||||
class TaskGroupRegister(MutableSet):
|
||||
"""
|
||||
This class combines the interface of a regular `set` with that of the `asyncio.Lock`.
|
||||
|
||||
It serves simultaneously as a container of IDs of tasks that belong to the same group, and as a mechanism for
|
||||
preventing race conditions within a task group. The lock should be acquired before cancelling the entire group of
|
||||
tasks, as well as before starting a task within the group.
|
||||
"""
|
||||
|
||||
def __init__(self, *task_ids: int) -> None:
|
||||
self._ids: Set[int] = set(task_ids)
|
||||
self._lock = Lock()
|
||||
|
||||
def __contains__(self, task_id: int) -> bool:
|
||||
"""Abstract method for the `MutableSet` base class."""
|
||||
return task_id in self._ids
|
||||
|
||||
def __iter__(self) -> Iterator[int]:
|
||||
"""Abstract method for the `MutableSet` base class."""
|
||||
return iter(self._ids)
|
||||
|
||||
def __len__(self) -> int:
|
||||
"""Abstract method for the `MutableSet` base class."""
|
||||
return len(self._ids)
|
||||
|
||||
def add(self, task_id: int) -> None:
|
||||
"""Abstract method for the `MutableSet` base class."""
|
||||
self._ids.add(task_id)
|
||||
|
||||
def discard(self, task_id: int) -> None:
|
||||
"""Abstract method for the `MutableSet` base class."""
|
||||
self._ids.discard(task_id)
|
||||
|
||||
async def acquire(self) -> bool:
|
||||
"""Wrapper around the lock's `acquire()` method."""
|
||||
return await self._lock.acquire()
|
||||
|
||||
def release(self) -> None:
|
||||
"""Wrapper around the lock's `release()` method."""
|
||||
self._lock.release()
|
||||
|
||||
async def __aenter__(self) -> None:
|
||||
"""Provides the asynchronous context manager syntax `async with ... :` when using the lock."""
|
||||
await self._lock.acquire()
|
||||
return None
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb) -> None:
|
||||
"""Provides the asynchronous context manager syntax `async with ... :` when using the lock."""
|
||||
self._lock.release()
|
Reference in New Issue
Block a user