krons 0.1.0__py3-none-any.whl

kronos/utils/concurrency/_async_call.py

@@ -0,0 +1,333 @@
+# Copyright (c) 2025 - 2026, HaiyangLi <quantocean.li at gmail dot com>
+# SPDX-License-Identifier: Apache-2.0
+
+"""Async batch processing with retry, timeout, and concurrency control.
+
+Primary exports:
+    alcall: Apply function to list elements concurrently with full control.
+    bcall: Batch processing wrapper yielding results per batch.
+"""
+
+from collections.abc import AsyncGenerator, Callable
+from typing import Any, ParamSpec, TypeVar
+
+from kronos.types._sentinel import Unset, not_sentinel
+from kronos.utils._lazy_init import LazyInit
+from kronos.utils._to_list import to_list
+
+from ._cancel import move_on_after
+from ._errors import get_cancelled_exc_class
+from ._patterns import non_cancel_subgroup
+from ._primitives import Semaphore
+from ._task import create_task_group
+from ._utils import is_coro_func, run_sync, sleep
+
+T = TypeVar("T")
+P = ParamSpec("P")
+
+_lazy = LazyInit()
+_MODEL_LIKE = None
+
+
+__all__ = (
+    "alcall",
+    "bcall",
+)
+
+
+def _do_init() -> None:
+    """Initialize Pydantic BaseModel detection for input normalization."""
+    global _MODEL_LIKE
+    from pydantic import BaseModel
+
+    _MODEL_LIKE = (BaseModel,)
+
+
+def _ensure_initialized() -> None:
+    """Trigger lazy Pydantic initialization on first use."""
+    _lazy.ensure(_do_init)
+
+
+def _validate_func(func: Any) -> Callable:
+    """Extract and validate a single callable.
+
+    Args:
+        func: Callable or single-element iterable containing a callable.
+
+    Returns:
+        The validated callable.
+
+    Raises:
+        ValueError: If not callable or iterable doesn't contain exactly one callable.
+    """
+    if callable(func):
+        return func
+
+    try:
+        func_list = list(func)
+    except TypeError:
+        raise ValueError("func must be callable or an iterable containing one callable.")
+
+    if len(func_list) != 1 or not callable(func_list[0]):
+        raise ValueError("Only one callable function is allowed.")
+
+    return func_list[0]
+
+
+def _normalize_input(
+    input_: Any,
+    *,
+    flatten: bool,
+    dropna: bool,
+    unique: bool,
+    flatten_tuple_set: bool,
+) -> list:
+    """Convert input to a flat list for batch processing.
+
+    Handles iterables, Pydantic models (as single items), and scalars.
+    """
+    if flatten or dropna:
+        return to_list(
+            input_,
+            flatten=flatten,
+            dropna=dropna,
+            unique=unique,
+            flatten_tuple_set=flatten_tuple_set,
+        )
+
+    if isinstance(input_, list):
+        return input_
+
+    if _MODEL_LIKE and isinstance(input_, _MODEL_LIKE):
+        return [input_]
+
+    try:
+        iter(input_)
+        return list(input_)
+    except TypeError:
+        return [input_]
+
+
+async def _call_with_timeout(
+    func: Callable,
+    item: Any,
+    is_coro: bool,
+    timeout: float | None,
+    **kwargs,
+) -> Any:
+    """Invoke function with optional timeout, handling both sync and async."""
+    if is_coro:
+        if timeout is not None:
+            with move_on_after(timeout) as cancel_scope:
+                result = await func(item, **kwargs)
+            if cancel_scope.cancelled_caught:
+                raise TimeoutError(f"Function call timed out after {timeout}s")
+            return result
+        return await func(item, **kwargs)
+    else:
+        if timeout is not None:
+            with move_on_after(timeout) as cancel_scope:
+                result = await run_sync(func, item, **kwargs)
+            if cancel_scope.cancelled_caught:
+                raise TimeoutError(f"Function call timed out after {timeout}s")
+            return result
+        return await run_sync(func, item, **kwargs)
+
+
+async def _execute_with_retry(
+    func: Callable,
+    item: Any,
+    index: int,
+    *,
+    is_coro: bool,
+    timeout: float | None,
+    initial_delay: float,
+    backoff: float,
+    max_attempts: int,
+    default: Any,
+    **kwargs,
+) -> tuple[int, Any]:
+    """Execute function with exponential backoff retry.
+
+    Returns (index, result) tuple to preserve ordering in concurrent execution.
+    Cancellation exceptions are never retried (respects structured concurrency).
+    """
+    attempts = 0
+    current_delay = initial_delay
+
+    while True:
+        try:
+            result = await _call_with_timeout(func, item, is_coro, timeout, **kwargs)
+            return index, result
+
+        except get_cancelled_exc_class():
+            raise
+
+        except Exception:
+            attempts += 1
+            if attempts <= max_attempts:
+                if current_delay:
+                    await sleep(current_delay)
+                    current_delay *= backoff
+            else:
+                if not_sentinel(default):
+                    return index, default
+                raise
+
+
+async def alcall(
+    input_: list[Any],
+    func: Callable[..., T],
+    /,
+    *,
+    input_flatten: bool = False,
+    input_dropna: bool = False,
+    input_unique: bool = False,
+    input_flatten_tuple_set: bool = False,
+    output_flatten: bool = False,
+    output_dropna: bool = False,
+    output_unique: bool = False,
+    output_flatten_tuple_set: bool = False,
+    delay_before_start: float = 0,
+    retry_initial_delay: float = 0,
+    retry_backoff: float = 1,
+    retry_default: Any = Unset,
+    retry_timeout: float | None = None,
+    retry_attempts: int = 0,
+    max_concurrent: int | None = None,
+    throttle_period: float | None = None,
+    return_exceptions: bool = False,
+    **kwargs: Any,
+) -> list[T | BaseException]:
+    """Apply function to each list element asynchronously with retry and concurrency control.
+
+    Args:
+        input_: List of items to process (or iterable that will be converted)
+        func: Callable to apply (sync or async)
+        input_flatten: Flatten nested input structures
+        input_dropna: Remove None/undefined from input
+        input_unique: Remove duplicate inputs (requires flatten)
+        input_flatten_tuple_set: Include tuples/sets in flattening
+        output_flatten: Flatten nested output structures
+        output_dropna: Remove None/undefined from output
+        output_unique: Remove duplicate outputs (requires flatten)
+        output_flatten_tuple_set: Include tuples/sets in output flattening
+        delay_before_start: Initial delay before processing (seconds)
+        retry_initial_delay: Initial retry delay (seconds)
+        retry_backoff: Backoff multiplier for retry delays
+        retry_default: Default value on retry exhaustion (Unset = raise)
+        retry_timeout: Timeout per function call (seconds)
+        retry_attempts: Maximum retry attempts (0 = no retry)
+        max_concurrent: Max concurrent executions (None = unlimited)
+        throttle_period: Delay between starting tasks (seconds)
+        return_exceptions: Return exceptions instead of raising
+        **kwargs: Additional arguments passed to func
+
+    Returns:
+        List of results (preserves input order, may include exceptions if return_exceptions=True)
+
+    Raises:
+        ValueError: If func is not callable
+        TimeoutError: If retry_timeout exceeded
+        ExceptionGroup: If return_exceptions=False and tasks raise
+    """
+    _ensure_initialized()
+
+    func = _validate_func(func)
+    input_ = _normalize_input(
+        input_,
+        flatten=input_flatten,
+        dropna=input_dropna,
+        unique=input_unique,
+        flatten_tuple_set=input_flatten_tuple_set,
+    )
+
+    if delay_before_start:
+        await sleep(delay_before_start)
+
+    semaphore = Semaphore(max_concurrent) if max_concurrent else None
+    throttle_delay = throttle_period or 0
+    is_coro = is_coro_func(func)
+    n_items = len(input_)
+    out: list[Any] = [None] * n_items
+
+    async def task_wrapper(item: Any, idx: int) -> None:
+        try:
+            if semaphore:
+                async with semaphore:
+                    _, result = await _execute_with_retry(
+                        func,
+                        item,
+                        idx,
+                        is_coro=is_coro,
+                        timeout=retry_timeout,
+                        initial_delay=retry_initial_delay,
+                        backoff=retry_backoff,
+                        max_attempts=retry_attempts,
+                        default=retry_default,
+                        **kwargs,
+                    )
+            else:
+                _, result = await _execute_with_retry(
+                    func,
+                    item,
+                    idx,
+                    is_coro=is_coro,
+                    timeout=retry_timeout,
+                    initial_delay=retry_initial_delay,
+                    backoff=retry_backoff,
+                    max_attempts=retry_attempts,
+                    default=retry_default,
+                    **kwargs,
+                )
+            out[idx] = result
+        except BaseException as exc:
+            out[idx] = exc
+            if not return_exceptions:
+                raise
+
+    try:
+        async with create_task_group() as tg:
+            for idx, item in enumerate(input_):
+                tg.start_soon(task_wrapper, item, idx)
+                if throttle_delay and idx < n_items - 1:
+                    await sleep(throttle_delay)
+    except ExceptionGroup as eg:
+        if not return_exceptions:
+            rest = non_cancel_subgroup(eg)
+            if rest is not None:
+                raise rest
+            raise
+
+    return to_list(
+        out,
+        flatten=output_flatten,
+        dropna=output_dropna,
+        unique=output_unique,
+        flatten_tuple_set=output_flatten_tuple_set,
+    )
+
+
+async def bcall(
+    input_: list[Any],
+    func: Callable[..., T],
+    /,
+    batch_size: int,
+    **kwargs: Any,
+) -> AsyncGenerator[list[T | BaseException], None]:
+    """Process input in batches using alcall. Yields results batch by batch.
+
+    Args:
+        input_: Items to process
+        func: Callable to apply
+        batch_size: Number of items per batch
+        **kwargs: Arguments passed to alcall (see alcall for details)
+
+    Yields:
+        List of results for each batch
+    """
+    input_ = to_list(input_, flatten=True, dropna=True)
+
+    for i in range(0, len(input_), batch_size):
+        batch = input_[i : i + batch_size]
+        yield await alcall(batch, func, **kwargs)

kronos/utils/concurrency/_cancel.py

@@ -0,0 +1,122 @@
+# Copyright (c) 2025 - 2026, HaiyangLi <quantocean.li at gmail dot com>
+# SPDX-License-Identifier: Apache-2.0
+
+"""Cancellation scope utilities wrapping anyio with None-safe timeouts.
+
+Provides context managers for timeout-based cancellation:
+- `fail_after`/`fail_at`: Raise TimeoutError on expiry
+- `move_on_after`/`move_on_at`: Silent cancellation on expiry
+
+All accept None to disable timeout while preserving outer scope cancellability.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterator
+from contextlib import contextmanager
+
+import anyio
+
+from ._utils import current_time
+
+CancelScope = anyio.CancelScope
+
+
+__all__ = (
+    "CancelScope",
+    "effective_deadline",
+    "fail_after",
+    "fail_at",
+    "move_on_after",
+    "move_on_at",
+)
+
+
+@contextmanager
+def fail_after(seconds: float | None) -> Iterator[CancelScope]:
+    """Context manager that raises TimeoutError after elapsed seconds.
+
+    Args:
+        seconds: Timeout duration, or None to disable timeout.
+
+    Yields:
+        CancelScope that can be checked via `scope.cancelled_caught`.
+    """
+    if seconds is None:
+        with CancelScope() as scope:
+            yield scope
+        return
+    with anyio.fail_after(seconds) as scope:
+        yield scope
+
+
+@contextmanager
+def move_on_after(seconds: float | None) -> Iterator[CancelScope]:
+    """Context manager that silently cancels after elapsed seconds.
+
+    Args:
+        seconds: Timeout duration, or None to disable timeout.
+
+    Yields:
+        CancelScope; check `scope.cancelled_caught` to detect timeout.
+    """
+    if seconds is None:
+        with CancelScope() as scope:
+            yield scope
+        return
+    with anyio.move_on_after(seconds) as scope:
+        yield scope
+
+
+@contextmanager
+def fail_at(deadline: float | None) -> Iterator[CancelScope]:
+    """Context manager that raises TimeoutError at absolute deadline.
+
+    Args:
+        deadline: Absolute time (from `current_time()`), or None to disable.
+
+    Yields:
+        CancelScope that can be checked via `scope.cancelled_caught`.
+    """
+    if deadline is None:
+        with CancelScope() as scope:
+            yield scope
+        return
+    now = current_time()
+    seconds = max(0.0, deadline - now)
+    with fail_after(seconds) as scope:
+        yield scope
+
+
+@contextmanager
+def move_on_at(deadline: float | None) -> Iterator[CancelScope]:
+    """Context manager that silently cancels at absolute deadline.
+
+    Args:
+        deadline: Absolute time (from `current_time()`), or None to disable.
+
+    Yields:
+        CancelScope; check `scope.cancelled_caught` to detect timeout.
+    """
+    if deadline is None:
+        with CancelScope() as scope:
+            yield scope
+        return
+    now = current_time()
+    seconds = max(0.0, deadline - now)
+    with anyio.move_on_after(seconds) as scope:
+        yield scope
+
+
+def effective_deadline() -> float | None:
+    """Return current effective deadline from enclosing cancel scopes.
+
+    Returns:
+        Absolute deadline time, -inf if already cancelled, or None if unlimited.
+
+    Note:
+        AnyIO uses +inf for "no deadline" and -inf for "already cancelled".
+        This function returns None for +inf but preserves -inf for detection.
+    """
+    d = anyio.current_effective_deadline()
+    return None if d == float("inf") else d

kronos/utils/concurrency/_errors.py

@@ -0,0 +1,96 @@
+# Copyright (c) 2025 - 2026, HaiyangLi <quantocean.li at gmail dot com>
+# SPDX-License-Identifier: Apache-2.0
+
+"""Cancellation error handling utilities.
+
+Provides backend-agnostic cancellation detection and ExceptionGroup helpers
+for cleanly separating cancellation from application errors.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Awaitable, Callable
+from typing import ParamSpec, TypeVar
+
+import anyio
+
+T = TypeVar("T")
+P = ParamSpec("P")
+
+
+__all__ = (
+    "get_cancelled_exc_class",
+    "is_cancelled",
+    "non_cancel_subgroup",
+    "shield",
+    "split_cancellation",
+)
+
+
+def get_cancelled_exc_class() -> type[BaseException]:
+    """Return backend-specific cancellation exception type.
+
+    Returns:
+        asyncio.CancelledError for asyncio, trio.Cancelled for trio.
+    """
+    return anyio.get_cancelled_exc_class()
+
+
+def is_cancelled(exc: BaseException) -> bool:
+    """Check if exception is a backend cancellation.
+
+    Args:
+        exc: Exception to check.
+
+    Returns:
+        True if exc is the backend's cancellation exception type.
+    """
+    return isinstance(exc, anyio.get_cancelled_exc_class())
+
+
+async def shield(func: Callable[P, Awaitable[T]], *args: P.args, **kwargs: P.kwargs) -> T:
+    """Execute async function protected from outer cancellation.
+
+    Args:
+        func: Async callable to shield.
+        *args: Positional arguments for func.
+        **kwargs: Keyword arguments for func.
+
+    Returns:
+        Result of func(*args, **kwargs).
+
+    Note:
+        Use sparingly. Shielded code cannot be cancelled, which may
+        delay shutdown. Prefer short critical sections only.
+    """
+    with anyio.CancelScope(shield=True):
+        result = await func(*args, **kwargs)
+    return result  # type: ignore[return-value]
+
+
+def split_cancellation(
+    eg: BaseExceptionGroup,
+) -> tuple[BaseExceptionGroup | None, BaseExceptionGroup | None]:
+    """Partition ExceptionGroup into cancellations and other errors.
+
+    Args:
+        eg: ExceptionGroup to split (Python 3.11+).
+
+    Returns:
+        Tuple of (cancellation_group, other_errors_group).
+        Either may be None if no matching exceptions.
+    """
+    return eg.split(anyio.get_cancelled_exc_class())
+
+
+def non_cancel_subgroup(eg: BaseExceptionGroup) -> BaseExceptionGroup | None:
+    """Extract non-cancellation exceptions from ExceptionGroup.
+
+    Args:
+        eg: ExceptionGroup to filter.
+
+    Returns:
+        ExceptionGroup of non-cancellation errors, or None if all were cancellations.
+    """
+    _, rest = eg.split(anyio.get_cancelled_exc_class())
+    return rest