fluidattacks_core_aio 12.0.0__py3-none-any.whl

fluidattacks_core/aio/__init__.py

@@ -0,0 +1,14 @@
+from fluidattacks_core.aio.processes import to_process
+from fluidattacks_core.aio.runners import run
+from fluidattacks_core.aio.tasks import as_completed, gather, merge_async_generators, to_background
+from fluidattacks_core.aio.threads import to_thread
+
+__all__ = [
+    "as_completed",
+    "gather",
+    "merge_async_generators",
+    "run",
+    "to_background",
+    "to_process",
+    "to_thread",
+]

fluidattacks_core/aio/processes.py

@@ -0,0 +1,64 @@
+import asyncio
+import contextvars
+import functools
+import threading
+from collections.abc import Callable
+from concurrent.futures import ProcessPoolExecutor
+from typing import ParamSpec, TypeVar
+
+P = ParamSpec("P")
+T = TypeVar("T")
+
+# Single-slot lazy holder, mutated in place so no ``global`` is needed. The
+# pool is created on first use (not at import), so importing this module has no
+# side effects and environments without ``/dev/shm`` (e.g. AWS Lambda) don't
+# crash at import.
+_POOL: dict[None, ProcessPoolExecutor] = {}
+_POOL_LOCK = threading.Lock()
+
+
+def _get_process_pool() -> ProcessPoolExecutor:
+    """Return the shared process pool, creating it on first use.
+
+    Environments without ``/dev/shm`` (e.g. AWS Lambda) cannot create the
+    pool's POSIX semaphores, so we fail explicitly and point the caller to the
+    thread-based alternative instead of silently degrading. Creation is guarded
+    by a lock so concurrent first callers (e.g. across event loops in separate
+    threads) share a single pool.
+    """
+    if None not in _POOL:
+        with _POOL_LOCK:
+            if None not in _POOL:
+                try:
+                    _POOL[None] = ProcessPoolExecutor()
+                except OSError as exc:
+                    msg = (
+                        "ProcessPoolExecutor is unavailable in this "
+                        "environment (no /dev/shm, e.g. AWS Lambda). "
+                        "Use to_thread instead."
+                    )
+                    raise RuntimeError(msg) from exc
+    return _POOL[None]
+
+
+# Public lazy accessor. ``PROCESS_POOL`` used to be an eagerly-created
+# ProcessPoolExecutor instance; it is now a ``() -> ProcessPoolExecutor``
+# callable built on first call. Callers use ``PROCESS_POOL()``. Changing the
+# type of this public entity is a breaking change, hence the major bump.
+PROCESS_POOL = _get_process_pool
+
+
+async def to_process(func: Callable[P, T], *args: P.args, **kwargs: P.kwargs) -> T:
+    """Asynchronously run a function in a separate process.
+
+    The underlying ProcessPoolExecutor sets a maximum number of worker processes
+    based on the number of cores. It is created lazily on first call; in
+    environments without ``/dev/shm`` (e.g. AWS Lambda) this raises a
+    ``RuntimeError`` pointing to ``to_thread``.
+
+    See: https://docs.python.org/3/library/concurrent.futures.html#concurrent.futures.ProcessPoolExecutor
+    """
+    loop = asyncio.get_running_loop()
+    ctx = contextvars.copy_context()
+    func_call = functools.partial(ctx.run, func, *args, **kwargs)
+    return await loop.run_in_executor(_get_process_pool(), func_call)

fluidattacks_core/aio/runners.py

@@ -0,0 +1,11 @@
+from collections.abc import Coroutine
+from typing import Any, TypeVar
+
+import uvloop
+
+T = TypeVar("T")
+
+
+def run(coroutine: Coroutine[Any, Any, T]) -> T:
+    """Run a coroutine and return the result, using uvloop."""
+    return uvloop.run(coroutine)

fluidattacks_core/aio/tasks.py

@@ -0,0 +1,213 @@
+import asyncio
+import logging
+from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Coroutine, Iterable, Iterator
+from contextlib import suppress
+from typing import Any, Literal, TypeVar, cast, overload
+
+T = TypeVar("T")
+
+# Sentinel object to signal generator completion
+_GENERATOR_DONE_SENTINEL = object()
+
+# Logger for the aio module
+_LOGGER = logging.getLogger(__name__)
+
+
+async def as_completed(
+    coroutines: Iterable[Awaitable[T]],
+    *,
+    concurrency_limit: int | None = None,
+) -> AsyncIterator[Awaitable[T]]:
+    """Run coroutines concurrently, yielding results in order of completion.
+
+    Args:
+        coroutines: An iterable of coroutines.
+        concurrency_limit: Maximum number of concurrent coroutines.
+
+    Yields:
+        Results from the coroutines in the order they complete.
+
+    """
+    if concurrency_limit:
+        semaphore = asyncio.Semaphore(concurrency_limit)
+
+        async def _run(coroutine: Awaitable[T]) -> T:
+            async with semaphore:
+                return await coroutine
+    else:
+
+        async def _run(coroutine: Awaitable[T]) -> T:
+            return await coroutine
+
+    tasks = [_run(coroutine) for coroutine in coroutines]
+    for task in asyncio.as_completed(tasks):
+        yield task
+
+
+@overload
+async def gather(
+    coroutines: Iterable[Awaitable[T]],
+    *,
+    concurrency_limit: int | None = None,
+    return_exceptions: Literal[False] = False,
+) -> list[T]: ...
+
+
+@overload
+async def gather(
+    coroutines: Iterable[Awaitable[T]],
+    *,
+    concurrency_limit: int | None = None,
+    return_exceptions: Literal[True],
+) -> list[T | BaseException]: ...
+
+
+async def gather(
+    coroutines: Iterable[Awaitable[T]],
+    *,
+    concurrency_limit: int | None = None,
+    return_exceptions: bool = False,
+) -> list[T] | list[T | BaseException]:
+    """Run coroutines concurrently.
+
+    Args:
+        coroutines: An iterable of coroutines.
+        concurrency_limit: Maximum number of concurrent coroutines.
+        return_exceptions: Whether to return exceptions instead of raising them.
+
+    Returns:
+        A list of results or exceptions.
+
+    Raises:
+        Exception: If return_exceptions is False and any coroutine raises an exception.
+
+    """
+    if concurrency_limit:
+        semaphore = asyncio.Semaphore(concurrency_limit)
+
+        async def _run(coroutine: Awaitable[T]) -> T:
+            async with semaphore:
+                return await coroutine
+    else:
+
+        async def _run(coroutine: Awaitable[T]) -> T:
+            return await coroutine
+
+    tasks = [_run(coroutine) for coroutine in coroutines]
+    return await asyncio.gather(*tasks, return_exceptions=return_exceptions)
+
+
+BACKGROUND_TASKS = set[asyncio.Task[Any]]()
+
+
+def to_background(coroutine: Coroutine[Any, Any, T]) -> None:
+    """Run a coroutine in the background, fire-and-forget style."""
+    task = asyncio.create_task(coroutine)
+    BACKGROUND_TASKS.add(task)
+    task.add_done_callback(BACKGROUND_TASKS.discard)
+
+
+async def _consume_generator(
+    gen: AsyncGenerator[T, None],
+    queue: asyncio.Queue[T | object],
+    active_generators: list[int],
+) -> None:
+    """Consume a single generator and put its items in the queue.
+
+    Args:
+        gen: The async generator to consume.
+        queue: The queue to put items into.
+        active_generators: A list containing the count of active generators.
+
+    """
+    active_generators[0] += 1
+    try:
+        async for item in gen:
+            await queue.put(item)
+    except Exception:
+        _LOGGER.exception(
+            "Error consuming generator %s, unhandled exception",
+            gen,
+        )
+    finally:
+        active_generators[0] -= 1
+        await queue.put(_GENERATOR_DONE_SENTINEL)
+
+
+def _start_initial_generator_tasks(
+    gen_iter: Iterator[AsyncGenerator[T, None]],
+    limit: int,
+    queue: asyncio.Queue[T | object],
+    active_generators: list[int],
+) -> list[asyncio.Task[None]]:
+    tasks: list[asyncio.Task[None]] = []
+    # Start initial tasks up to the limit
+    for _ in range(limit):
+        try:
+            gen = next(gen_iter)
+            tasks.append(asyncio.create_task(_consume_generator(gen, queue, active_generators)))
+        except StopIteration:
+            break
+    return tasks
+
+
+async def merge_async_generators(
+    generators: Iterable[AsyncGenerator[T, None]],
+    limit: int,
+) -> AsyncGenerator[T, None]:
+    """Merge multiple async generators into a single async generator with concurrency control.
+
+    This function runs multiple async generators concurrently, yielding items as they become
+    available. It maintains a maximum number of active generators specified by the limit.
+
+    Args:
+        generators: An iterable of async generators to merge.
+        limit: Maximum number of generators to run concurrently.
+
+    Yields:
+        Items from the generators as they become available.
+
+    Raises:
+        ValueError: If limit is less than 1.
+
+    Example:
+        >>> async def gen1():
+        ...     yield 1
+        ...     yield 2
+        >>> async def gen2():
+        ...     yield 3
+        ...     yield 4
+        >>> async for item in merge_async_generators([gen1(), gen2()], limit=2):
+        ...     print(item)  # May print 1, 3, 2, 4 in any order
+
+    """
+    if limit < 1:
+        msg = "limit must be at least 1"
+        raise ValueError(msg)
+
+    queue: asyncio.Queue[T | object] = asyncio.Queue()
+    active_generators = [0]  # Use list to allow modification in nested function
+    gen_iter = iter(generators)
+    tasks = _start_initial_generator_tasks(gen_iter, limit, queue, active_generators)
+    # Keep track of how many generators we expect to finish
+    expected_done_signals = len(tasks)
+    done_signals_received = 0
+
+    try:
+        while done_signals_received < expected_done_signals:
+            item = await queue.get()
+            if item is _GENERATOR_DONE_SENTINEL:
+                done_signals_received += 1
+                # Try to start a new task if there are more generators
+                with suppress(StopIteration):
+                    next_gen = next(gen_iter)
+                    tasks.append(
+                        asyncio.create_task(_consume_generator(next_gen, queue, active_generators)),
+                    )
+                    expected_done_signals += 1
+            else:
+                yield cast("T", item)
+    finally:
+        # Ensure all tasks are complete
+        if tasks:
+            await asyncio.gather(*tasks, return_exceptions=True)

fluidattacks_core/aio/threads.py

@@ -0,0 +1,17 @@
+import asyncio
+from collections.abc import Callable
+from typing import ParamSpec, TypeVar
+
+P = ParamSpec("P")
+T = TypeVar("T")
+
+
+async def to_thread(func: Callable[P, T], *args: P.args, **kwargs: P.kwargs) -> T:
+    """Asynchronously run a function in a separate thread.
+
+    The underlying ThreadPoolExecutor sets a maximum number of worker threads
+    based on the number of cores.
+
+    See: https://docs.python.org/3/library/concurrent.futures.html#concurrent.futures.ThreadPoolExecutor
+    """
+    return await asyncio.to_thread(func, *args, **kwargs)

fluidattacks_core_aio-12.0.0.dist-info/METADATA

@@ -0,0 +1,14 @@
+Metadata-Version: 2.5
+Name: fluidattacks_core_aio
+Version: 12.0.0
+Summary: Fluid Attacks Core Aio Library
+Author-email: Development <development@fluidattacks.com>
+License: MPL-2.0
+Classifier: Development Status :: 1 - Planning
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: Mozilla Public License 2.0 (MPL 2.0)
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Topic :: Software Development :: Libraries
+Requires-Python: >=3.11
+Requires-Dist: uvloop>=0.21.0

fluidattacks_core_aio-12.0.0.dist-info/RECORD

@@ -0,0 +1,9 @@
+fluidattacks_core/aio/__init__.py,sha256=Xn4sNbAYtqLPV-em8vWVQD5eRH3IigDpAVzHnky0onY,399
+fluidattacks_core/aio/processes.py,sha256=E-Jw2h_8eOQIWN8AeyVAt_2w7Lv1E-F7isz_Ox6g0HY,2620
+fluidattacks_core/aio/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+fluidattacks_core/aio/runners.py,sha256=rJYtrshcWawhWWWagRn6Rsre_BsmAECLWG0FYOheLJQ,250
+fluidattacks_core/aio/tasks.py,sha256=bZ-d0BP4ptCK9seQiXY4nFnFGGK3ZhIDQTM52k-yMDc,6487
+fluidattacks_core/aio/threads.py,sha256=JrnAv7_jAlDDj2AI5MviWvxU8EHQVpW-c7uCtst-2yY,556
+fluidattacks_core_aio-12.0.0.dist-info/METADATA,sha256=ebk5Djl8u5M68jQqJqGObaXayOdnErWJ4wAoMm1Bh-4,555
+fluidattacks_core_aio-12.0.0.dist-info/WHEEL,sha256=zOwg4jB6zX2kU910N-cMawjivD6tO8NEWvE12je1bVk,87
+fluidattacks_core_aio-12.0.0.dist-info/RECORD,,

fluidattacks_core_aio-12.0.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.32.0
+Root-Is-Purelib: true
+Tag: py3-none-any