haiway 0.10.15__py3-none-any.whl → 0.10.17__py3-none-any.whl

haiway/helpers/throttling.py

@@ -0,0 +1,133 @@
+from asyncio import (
+    Lock,
+    iscoroutinefunction,
+    sleep,
+)
+from collections import deque
+from collections.abc import Callable, Coroutine
+from datetime import timedelta
+from time import monotonic
+from typing import cast, overload
+
+from haiway.utils.mimic import mimic_function
+
+__all__ = [
+    "throttle",
+]
+
+
+@overload
+def throttle[**Args, Result](
+    function: Callable[Args, Coroutine[None, None, Result]],
+    /,
+) -> Callable[Args, Coroutine[None, None, Result]]: ...
+
+
+@overload
+def throttle[**Args, Result](
+    *,
+    limit: int = 1,
+    period: timedelta | float = 1,
+) -> Callable[
+    [Callable[Args, Coroutine[None, None, Result]]], Callable[Args, Coroutine[None, None, Result]]
+]: ...
+
+
+def throttle[**Args, Result](
+    function: Callable[Args, Coroutine[None, None, Result]] | None = None,
+    *,
+    limit: int = 1,
+    period: timedelta | float = 1,
+) -> (
+    Callable[
+        [Callable[Args, Coroutine[None, None, Result]]],
+        Callable[Args, Coroutine[None, None, Result]],
+    ]
+    | Callable[Args, Coroutine[None, None, Result]]
+):
+    """\
+    Throttle for function calls with custom limit and period time. \
+    Works only for async functions by waiting desired time before execution. \
+    It is not allowed to be used on class or instance methods. \
+    This wrapper is not thread safe.
+
+    Parameters
+    ----------
+    function: Callable[Args, Coroutine[None, None, Result]]
+        function to wrap in throttle
+    limit: int
+        limit of executions in given period, if no period was specified
+        it is number of concurrent executions instead, default is 1
+    period: timedelta | float | None
+        period time (in seconds by default) during which the limit resets, default is 1 second
+
+    Returns
+    -------
+    Callable[[Callable[Args, Coroutine[None, None, Result]]], Callable[Args, Coroutine[None, None, Result]]] \
+    | Callable[Args, Coroutine[None, None, Result]]
+        provided function wrapped in throttle
+    """  # noqa: E501
+
+    def _wrap(
+        function: Callable[Args, Coroutine[None, None, Result]],
+    ) -> Callable[Args, Coroutine[None, None, Result]]:
+        assert iscoroutinefunction(function)  # nosec: B101
+        return cast(
+            Callable[Args, Coroutine[None, None, Result]],
+            _AsyncThrottle(
+                function,
+                limit=limit,
+                period=period,
+            ),
+        )
+
+    if function := function:
+        return _wrap(function)
+
+    else:
+        return _wrap
+
+
+class _AsyncThrottle[**Args, Result]:
+    def __init__(
+        self,
+        function: Callable[Args, Coroutine[None, None, Result]],
+        /,
+        limit: int,
+        period: timedelta | float,
+    ) -> None:
+        self._function: Callable[Args, Coroutine[None, None, Result]] = function
+        self._entries: deque[float] = deque()
+        self._lock: Lock = Lock()
+        self._limit: int = limit
+        self._period: float
+        match period:
+            case timedelta() as delta:
+                self._period = delta.total_seconds()
+
+            case period_seconds:
+                self._period = period_seconds
+
+        # mimic function attributes if able
+        mimic_function(function, within=self)
+
+    async def __call__(
+        self,
+        *args: Args.args,
+        **kwargs: Args.kwargs,
+    ) -> Result:
+        async with self._lock:
+            time_now: float = monotonic()
+            while self._entries:  # cleanup old entries
+                if self._entries[0] + self._period <= time_now:
+                    self._entries.popleft()
+
+                else:
+                    break
+
+            if len(self._entries) >= self._limit:
+                await sleep(self._entries[0] - time_now)
+
+            self._entries.append(monotonic())
+
+        return await self._function(*args, **kwargs)

haiway/helpers/timeouted.py

@@ -0,0 +1,112 @@
+from asyncio import AbstractEventLoop, Future, Task, TimerHandle, get_running_loop
+from collections.abc import Callable, Coroutine
+
+from haiway.utils.mimic import mimic_function
+
+__all__ = [
+    "timeout",
+]
+
+
+def timeout[**Args, Result](
+    timeout: float,
+    /,
+) -> Callable[
+    [Callable[Args, Coroutine[None, None, Result]]],
+    Callable[Args, Coroutine[None, None, Result]],
+]:
+    """\
+    Timeout wrapper for a function call. \
+    When the timeout time will pass before function returns function execution will be \
+    cancelled and TimeoutError exception will raise. Make sure that wrapped \
+    function handles cancellation properly.
+    This wrapper is not thread safe.
+
+    Parameters
+    ----------
+    timeout: float
+        timeout time in seconds
+
+    Returns
+    -------
+    Callable[[Callable[_Args, _Result]], Callable[_Args, _Result]] | Callable[_Args, _Result]
+        function wrapper adding timeout
+    """
+
+    def _wrap(
+        function: Callable[Args, Coroutine[None, None, Result]],
+    ) -> Callable[Args, Coroutine[None, None, Result]]:
+        return _AsyncTimeout(
+            function,
+            timeout=timeout,
+        )
+
+    return _wrap
+
+
+class _AsyncTimeout[**Args, Result]:
+    def __init__(
+        self,
+        function: Callable[Args, Coroutine[None, None, Result]],
+        /,
+        timeout: float,
+    ) -> None:
+        self._function: Callable[Args, Coroutine[None, None, Result]] = function
+        self._timeout: float = timeout
+
+        # mimic function attributes if able
+        mimic_function(function, within=self)
+
+    async def __call__(
+        self,
+        *args: Args.args,
+        **kwargs: Args.kwargs,
+    ) -> Result:
+        loop: AbstractEventLoop = get_running_loop()
+        future: Future[Result] = loop.create_future()
+        task: Task[Result] = loop.create_task(
+            self._function(
+                *args,
+                **kwargs,
+            ),
+        )
+
+        def on_timeout(
+            future: Future[Result],
+        ) -> None:
+            if future.done():
+                return  # ignore if already finished
+
+            # result future on its completion will ensure that task will complete
+            future.set_exception(TimeoutError())
+
+        timeout_handle: TimerHandle = loop.call_later(
+            self._timeout,
+            on_timeout,
+            future,
+        )
+
+        def on_completion(
+            task: Task[Result],
+        ) -> None:
+            timeout_handle.cancel()  # at this stage we no longer need timeout to trigger
+
+            if future.done():
+                return  # ignore if already finished
+
+            try:
+                future.set_result(task.result())
+
+            except Exception as exc:
+                future.set_exception(exc)
+
+        task.add_done_callback(on_completion)
+
+        def on_result(
+            future: Future[Result],
+        ) -> None:
+            task.cancel()  # when result future completes make sure that task also completes
+
+        future.add_done_callback(on_result)
+
+        return await future

haiway/helpers/tracing.py

@@ -0,0 +1,137 @@
+from asyncio import iscoroutinefunction
+from collections.abc import Callable, Coroutine, Mapping, Sequence
+from typing import Any, Self, cast
+
+from haiway.context import ctx
+from haiway.state import State
+from haiway.types import MISSING, Missing
+from haiway.utils import mimic_function
+
+__all__ = [
+    "ArgumentsTrace",
+    "ResultTrace",
+    "traced",
+]
+
+
+class ArgumentsTrace(State):
+    if __debug__:
+
+        @classmethod
+        def of(cls, *args: Any, **kwargs: Any) -> Self:
+            return cls(
+                args=args if args else MISSING,
+                kwargs=kwargs if kwargs else MISSING,
+            )
+
+    else:  # remove tracing for non debug runs to prevent accidental secret leaks
+
+        @classmethod
+        def of(cls, *args: Any, **kwargs: Any) -> Self:
+            return cls(
+                args=MISSING,
+                kwargs=MISSING,
+            )
+
+    args: Sequence[Any] | Missing
+    kwargs: Mapping[str, Any] | Missing
+
+
+class ResultTrace(State):
+    if __debug__:
+
+        @classmethod
+        def of(
+            cls,
+            value: Any,
+            /,
+        ) -> Self:
+            return cls(result=value)
+
+    else:  # remove tracing for non debug runs to prevent accidental secret leaks
+
+        @classmethod
+        def of(
+            cls,
+            value: Any,
+            /,
+        ) -> Self:
+            return cls(result=MISSING)
+
+    result: Any | Missing
+
+
+def traced[**Args, Result](
+    function: Callable[Args, Result],
+    /,
+) -> Callable[Args, Result]:
+    if __debug__:
+        if iscoroutinefunction(function):
+            return cast(
+                Callable[Args, Result],
+                _traced_async(
+                    function,
+                    label=function.__name__,
+                ),
+            )
+
+        else:
+            return _traced_sync(
+                function,
+                label=function.__name__,
+            )
+
+    else:  # do not trace on non debug runs
+        return function
+
+
+def _traced_sync[**Args, Result](
+    function: Callable[Args, Result],
+    /,
+    label: str,
+) -> Callable[Args, Result]:
+    def traced(
+        *args: Args.args,
+        **kwargs: Args.kwargs,
+    ) -> Result:
+        with ctx.scope(label):
+            ctx.record(ArgumentsTrace.of(*args, **kwargs))
+            try:
+                result: Result = function(*args, **kwargs)
+                ctx.record(ResultTrace.of(result))
+                return result
+
+            except BaseException as exc:
+                ctx.record(ResultTrace.of(f"{type(exc)}: {exc}"))
+                raise exc
+
+    return mimic_function(
+        function,
+        within=traced,
+    )
+
+
+def _traced_async[**Args, Result](
+    function: Callable[Args, Coroutine[Any, Any, Result]],
+    /,
+    label: str,
+) -> Callable[Args, Coroutine[Any, Any, Result]]:
+    async def traced(
+        *args: Args.args,
+        **kwargs: Args.kwargs,
+    ) -> Result:
+        with ctx.scope(label):
+            ctx.record(ArgumentsTrace.of(*args, **kwargs))
+            try:
+                result: Result = await function(*args, **kwargs)
+                ctx.record(ResultTrace.of(result))
+                return result
+
+            except BaseException as exc:
+                ctx.record(ResultTrace.of(f"{type(exc)}: {exc}"))
+                raise exc
+
+    return mimic_function(
+        function,
+        within=traced,
+    )

haiway/state/__init__.py

@@ -0,0 +1,12 @@
+from haiway.state.attributes import AttributeAnnotation, attribute_annotations
+from haiway.state.path import AttributePath
+from haiway.state.requirement import AttributeRequirement
+from haiway.state.structure import State
+
+__all__ = [
+    "AttributeAnnotation",
+    "AttributePath",
+    "AttributeRequirement",
+    "State",
+    "attribute_annotations",
+]