coco-framework 0.1.0__py3-none-any.whl

coco_core/utils/aiter.py

@@ -0,0 +1,357 @@
+"""Asynchronous iterator utilities.
+
+Adapted from
+https://github.com/maxfischer2781/asyncstdlib/blob/master/asyncstdlib/itertools.py
+MIT License.
+"""
+
+from collections import deque
+from collections.abc import (
+    AsyncGenerator,
+    AsyncIterable,
+    AsyncIterator,
+    Awaitable,
+    Callable,
+    Iterator,
+)
+from contextlib import AbstractAsyncContextManager
+from types import TracebackType
+from typing import (
+    Any,
+    Generic,
+    TypeVar,
+    cast,
+    overload,
+)
+
+from typing_extensions import override
+
+from coco_core._api.deprecation import deprecated
+
+T = TypeVar("T")
+
+_no_default = object()
+
+
+# https://github.com/python/cpython/blob/main/Lib/test/test_asyncgen.py#L54
+@deprecated(since="1.1.2", removal="2.0.0")
+def py_anext(
+    iterator: AsyncIterator[T], default: T | Any = _no_default
+) -> Awaitable[T | Any | None]:
+    """Pure-Python implementation of `anext()` for testing purposes.
+
+    Closely matches the builtin `anext()` C implementation.
+
+    Can be used to compare the built-in implementation of the inner coroutines machinery
+    to C-implementation of `__anext__()` and `send()` or `throw()` on the returned
+    generator.
+
+    Args:
+        iterator: The async iterator to advance.
+        default: The value to return if the iterator is exhausted.
+
+            If not provided, a `StopAsyncIteration` exception is raised.
+
+    Returns:
+        The next value from the iterator, or the default value if the iterator is
+            exhausted.
+
+    Raises:
+        TypeError: If the iterator is not an async iterator.
+    """
+    try:
+        __anext__ = cast(
+            "Callable[[AsyncIterator[T]], Awaitable[T]]", type(iterator).__anext__
+        )
+    except AttributeError as e:
+        msg = f"{iterator!r} is not an async iterator"
+        raise TypeError(msg) from e
+
+    if default is _no_default:
+        return __anext__(iterator)
+
+    async def anext_impl() -> T | Any:
+        try:
+            # The C code is way more low-level than this, as it implements
+            # all methods of the iterator protocol. In this implementation
+            # we're relying on higher-level coroutine concepts, but that's
+            # exactly what we want -- crosstest pure-Python high-level
+            # implementation and low-level C anext() iterators.
+            return await __anext__(iterator)
+        except StopAsyncIteration:
+            return default
+
+    return anext_impl()
+
+
+class NoLock:
+    """Dummy lock that provides the proper interface but no protection."""
+
+    async def __aenter__(self) -> None:
+        """Do nothing."""
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> bool:
+        """Return False, exception not suppressed."""
+        return False
+
+
+async def tee_peer(
+    iterator: AsyncIterator[T],
+    # the buffer specific to this peer
+    buffer: deque[T],
+    # the buffers of all peers, including our own
+    peers: list[deque[T]],
+    lock: AbstractAsyncContextManager[Any],
+) -> AsyncGenerator[T, None]:
+    """An individual iterator of a `tee`.
+
+    This function is a generator that yields items from the shared iterator
+    `iterator`. It buffers items until the least advanced iterator has yielded them as
+    well.
+
+    The buffer is shared with all other peers.
+
+    Args:
+        iterator: The shared iterator.
+        buffer: The buffer for this peer.
+        peers: The buffers of all peers.
+        lock: The lock to synchronise access to the shared buffers.
+
+    Yields:
+        The next item from the shared iterator.
+    """
+    try:
+        while True:
+            if not buffer:
+                async with lock:
+                    # Another peer produced an item while we were waiting for the lock.
+                    # Proceed with the next loop iteration to yield the item.
+                    if buffer:
+                        continue
+                    try:
+                        item = await anext(iterator)
+                    except StopAsyncIteration:
+                        break
+                    else:
+                        # Append to all buffers, including our own. We'll fetch our
+                        # item from the buffer again, instead of yielding it directly.
+                        # This ensures the proper item ordering if any of our peers
+                        # are fetching items concurrently. They may have buffered their
+                        # item already.
+                        for peer_buffer in peers:
+                            peer_buffer.append(item)
+            yield buffer.popleft()
+    finally:
+        async with lock:
+            # this peer is done - remove its buffer
+            for idx, peer_buffer in enumerate(peers):  # pragma: no branch
+                if peer_buffer is buffer:
+                    peers.pop(idx)
+                    break
+            # if we are the last peer, try and close the iterator
+            if not peers and hasattr(iterator, "aclose"):
+                await iterator.aclose()
+
+
+class Tee(Generic[T]):
+    """Create `n` separate asynchronous iterators over `iterable`.
+
+    This splits a single `iterable` into multiple iterators, each providing
+    the same items in the same order.
+
+    All child iterators may advance separately but share the same items from `iterable`
+    -- when the most advanced iterator retrieves an item, it is buffered until the least
+    advanced iterator has yielded it as well.
+
+    A `tee` works lazily and can handle an infinite `iterable`, provided
+    that all iterators advance.
+
+    ```python
+    async def derivative(sensor_data):
+        previous, current = a.tee(sensor_data, n=2)
+        await a.anext(previous)  # advance one iterator
+        return a.map(operator.sub, previous, current)
+    ```
+
+    Unlike `itertools.tee`, `.tee` returns a custom type instead of a `tuple`. Like a
+    tuple, it can be indexed, iterated and unpacked to get the child iterators. In
+    addition, its `.tee.aclose` method immediately closes all children, and it can be
+    used in an `async with` context for the same effect.
+
+    If `iterable` is an iterator and read elsewhere, `tee` will *not* provide these
+    items. Also, `tee` must internally buffer each item until the last iterator has
+    yielded it; if the most and least advanced iterator differ by most data, using a
+    `list` is more efficient (but not lazy).
+
+    If the underlying iterable is concurrency safe (`anext` may be awaited concurrently)
+    the resulting iterators are concurrency safe as well. Otherwise, the iterators are
+    safe if there is only ever one single "most advanced" iterator.
+
+    To enforce sequential use of `anext`, provide a `lock`
+
+    - e.g. an `asyncio.Lock` instance in an `asyncio` application - and access is
+        automatically synchronised.
+
+    """
+
+    def __init__(
+        self,
+        iterable: AsyncIterator[T],
+        n: int = 2,
+        *,
+        lock: AbstractAsyncContextManager[Any] | None = None,
+    ):
+        """Create a `tee`.
+
+        Args:
+            iterable: The iterable to split.
+            n: The number of iterators to create.
+            lock: The lock to synchronise access to the shared buffers.
+
+        """
+        self._iterator = aiter(iterable)
+        self._buffers: list[deque[T]] = [deque() for _ in range(n)]
+        self._children = tuple(
+            tee_peer(
+                iterator=self._iterator,
+                buffer=buffer,
+                peers=self._buffers,
+                lock=lock if lock is not None else NoLock(),
+            )
+            for buffer in self._buffers
+        )
+
+    def __len__(self) -> int:
+        """Return the number of child iterators."""
+        return len(self._children)
+
+    @overload
+    def __getitem__(self, item: int) -> AsyncIterator[T]: ...
+
+    @overload
+    def __getitem__(self, item: slice) -> tuple[AsyncIterator[T], ...]: ...
+
+    def __getitem__(
+        self, item: int | slice
+    ) -> AsyncIterator[T] | tuple[AsyncIterator[T], ...]:
+        """Return the child iterator(s) for the given index or slice."""
+        return self._children[item]
+
+    def __iter__(self) -> Iterator[AsyncIterator[T]]:
+        """Iterate over the child iterators.
+
+        Yields:
+            The child iterators.
+        """
+        yield from self._children
+
+    async def __aenter__(self) -> "Tee[T]":
+        """Return the tee instance."""
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> bool:
+        """Close all child iterators.
+
+        Returns:
+            `False`, exceptions not suppressed.
+        """
+        await self.aclose()
+        return False
+
+    async def aclose(self) -> None:
+        """Async close all child iterators."""
+        for child in self._children:
+            await child.aclose()
+
+
+atee = Tee
+
+
+class aclosing(AbstractAsyncContextManager[Any]):  # noqa: N801
+    """Async context manager to wrap an `AsyncGenerator` that has a `aclose()` method.
+
+    Code like this:
+
+    ```python
+    async with aclosing(<module>.fetch(<arguments>)) as agen:
+        <block>
+    ```
+
+    ...is equivalent to this:
+
+    ```python
+    agen = <module>.fetch(<arguments>)
+    try:
+        <block>
+    finally:
+        await agen.aclose()
+
+    ```
+    """
+
+    def __init__(self, thing: AsyncGenerator[Any, Any] | AsyncIterator[Any]) -> None:
+        """Create the context manager.
+
+        Args:
+            thing: The resource to wrap.
+        """
+        self.thing = thing
+
+    @override
+    async def __aenter__(self) -> AsyncGenerator[Any, Any] | AsyncIterator[Any]:
+        return self.thing
+
+    @override
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        if hasattr(self.thing, "aclose"):
+            await self.thing.aclose()
+
+
+async def abatch_iterate(
+    size: int | None, iterable: AsyncIterable[T]
+) -> AsyncIterator[list[T]]:
+    """Utility batching function for async iterables.
+
+    Args:
+        size: The size of the batch.
+
+            If `None`, returns a single batch.
+        iterable: The async iterable to batch.
+
+    Yields:
+        The batches.
+
+    Raises:
+        ValueError: If `size` is not `None` and is not a positive integer.
+    """
+    if size is None:
+        single_batch = [el async for el in iterable]
+        if single_batch:
+            yield single_batch
+        return
+    if size <= 0:
+        msg = f"Batch size must be a positive integer, got {size}."
+        raise ValueError(msg)
+    batch: list[T] = []
+    async for element in iterable:
+        batch.append(element)
+        if len(batch) >= size:
+            yield batch
+            batch = []
+    if batch:
+        yield batch

coco_core/utils/env.py

@@ -0,0 +1,86 @@
+"""Utilities for environment variables."""
+
+from __future__ import annotations
+
+import os
+from typing import Any
+
+
+def env_var_is_set(env_var: str) -> bool:
+    """Check if an environment variable is set.
+
+    Args:
+        env_var: The name of the environment variable.
+
+    Returns:
+        `True` if the environment variable is set, `False` otherwise.
+    """
+    return env_var in os.environ and os.environ[env_var] not in {
+        "",
+        "0",
+        "false",
+        "False",
+    }
+
+
+def get_from_dict_or_env(
+    data: dict[str, Any],
+    key: str | list[str],
+    env_key: str,
+    default: str | None = None,
+) -> str:
+    """Get a value from a dictionary or an environment variable.
+
+    Args:
+        data: The dictionary to look up the key in.
+        key: The key to look up in the dictionary.
+
+            This can be a list of keys to try in order.
+        env_key: The environment variable to look up if the key is not
+            in the dictionary.
+        default: The default value to return if the key is not in the dictionary
+            or the environment.
+
+    Returns:
+        The dict value or the environment variable value.
+    """
+    if isinstance(key, (list, tuple)):
+        for k in key:
+            if value := data.get(k):
+                return str(value)
+
+    if isinstance(key, str) and key in data and data[key]:
+        return str(data[key])
+
+    key_for_err = key[0] if isinstance(key, (list, tuple)) else key
+
+    return get_from_env(key_for_err, env_key, default=default)
+
+
+def get_from_env(key: str, env_key: str, default: str | None = None) -> str:
+    """Get a value from a dictionary or an environment variable.
+
+    Args:
+        key: The key to look up in the dictionary.
+        env_key: The environment variable to look up if the key is not
+            in the dictionary.
+        default: The default value to return if the key is not in the dictionary
+            or the environment.
+
+    Returns:
+        The value of the key.
+
+    Raises:
+        ValueError: If the key is not in the dictionary and no default value is
+            provided or if the environment variable is not set.
+    """
+    if env_value := os.getenv(env_key):
+        return env_value
+    if default is not None:
+        return default
+    msg = (
+        f"Did not find {key}, please add an environment variable"
+        f" `{env_key}` which contains it, or pass"
+        f" `{key}` as a named parameter."
+    )
+    raise ValueError(msg)

coco_core/utils/formatting.py

@@ -0,0 +1,80 @@
+"""Utilities for formatting strings."""
+
+from collections.abc import Mapping, Sequence
+from string import Formatter
+from typing import Any
+
+
+class StrictFormatter(Formatter):
+    """A string formatter that enforces keyword-only argument substitution.
+
+    This formatter extends Python's built-in `string.Formatter` to provide stricter
+    validation for prompt template formatting. It ensures that all variable
+    substitutions use keyword arguments rather than positional arguments, which improves
+    clarity and reduces errors when formatting prompt templates.
+
+    Example:
+        >>> fmt = StrictFormatter()
+        >>> fmt.format("Hello, {name}!", name="World")
+        'Hello, World!'
+        >>> fmt.format("Hello, {}!", "World")  # Raises ValueError
+    """
+
+    def vformat(
+        self, format_string: str, args: Sequence[Any], kwargs: Mapping[str, Any]
+    ) -> str:
+        """Format a string using only keyword arguments.
+
+        Overrides the base `vformat` to reject positional arguments, ensuring all
+        substitutions are explicit and named.
+
+        Args:
+            format_string: A string containing replacement fields (e.g., `'{name}'`).
+            args: Positional arguments (must be empty).
+            kwargs: Keyword arguments for substitution into the format string.
+
+        Returns:
+            The formatted string with all replacement fields substituted.
+
+        Raises:
+            ValueError: If any positional arguments are provided.
+        """
+        if len(args) > 0:
+            msg = (
+                "No arguments should be provided, "
+                "everything should be passed as keyword arguments."
+            )
+            raise ValueError(msg)
+        return super().vformat(format_string, args, kwargs)
+
+    def validate_input_variables(
+        self, format_string: str, input_variables: list[str]
+    ) -> None:
+        """Validate that input variables match the placeholders in a format string.
+
+        Checks that the provided input variables can be used to format the given string
+        without missing or extra keys. This is useful for validating prompt templates
+        before runtime.
+
+        Args:
+            format_string: A string containing replacement fields to validate
+                against (e.g., `'Hello, {name}!'`).
+            input_variables: List of variable names expected to fill the
+                replacement fields.
+
+        Raises:
+            KeyError: If the format string contains placeholders not present
+                in input_variables.
+
+        Example:
+            >>> fmt = StrictFormatter()
+            >>> fmt.validate_input_variables("Hello, {name}!", ["name"])  # OK
+            >>> fmt.validate_input_variables("Hello, {name}!", ["other"])  # Raises
+        """
+        dummy_inputs = dict.fromkeys(input_variables, "foo")
+        super().format(format_string, **dummy_inputs)
+
+
+#: Default StrictFormatter instance for use throughout Coco.
+#: Used internally for formatting prompt templates with named variables.
+formatter = StrictFormatter()