[pull] main from apache:main

burr/common/async_utils.py

@@ -16,7 +16,7 @@
 # under the License.
 
 import inspect
-from typing import AsyncGenerator, AsyncIterable, Generator, List, TypeVar, Union
+from typing import Any, AsyncGenerator, AsyncIterable, Coroutine, Generator, List, TypeVar, Union
 
 T = TypeVar("T")
 
@@ -27,6 +27,46 @@
 SyncOrAsyncGeneratorOrItemOrList = Union[SyncOrAsyncGenerator[GenType], List[GenType], GenType]
 
 
+class _AsyncPersisterContextManager:
+    """Wraps an async coroutine that returns a persister so it can be used
+    directly with ``async with``::
+
+        async with AsyncSQLitePersister.from_values(...) as persister:
+            ...
+
+    The wrapper awaits the coroutine on ``__aenter__`` and delegates
+    ``__aexit__`` to the persister's own ``__aexit__``.
+
+    .. note::
+        Each instance wraps a single coroutine and can only be consumed once,
+        either via ``await`` or ``async with``. A second use will raise
+        ``RuntimeError``.
+    """
+
+    def __init__(self, coro: Coroutine[Any, Any, Any]):
+        self._coro = coro
+        self._persister = None
+        self._consumed = False
+
+    def __await__(self):
+        if self._consumed:
+            raise RuntimeError("This factory result has already been consumed")
+        self._consumed = True
+        return self._coro.__await__()
+
+    async def __aenter__(self):
+        if self._consumed:
+            raise RuntimeError("This factory result has already been consumed")
+        self._consumed = True
+        self._persister = await self._coro
+        return await self._persister.__aenter__()
+
+    async def __aexit__(self, exc_type, exc_value, traceback):
+        if self._persister is None:
+            return False
+        return await self._persister.__aexit__(exc_type, exc_value, traceback)
+
+
 async def asyncify_generator(
     generator: SyncOrAsyncGenerator[GenType],
 ) -> AsyncGenerator[GenType, None]:

burr/core/__init__.py

@@ -15,7 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
-from burr.core.action import Action, Condition, Result, action, default, expr, when
+from burr.core.action import Action, Condition, Result, action, default, expr, type_eraser, when
 from burr.core.application import (
     Application,
     ApplicationBuilder,
@@ -35,6 +35,7 @@
     "Condition",
     "default",
     "expr",
+    "type_eraser",
     "Result",
     "State",
     "when",

burr/core/action.py

@@ -100,8 +100,75 @@ def visit_Subscript(self, node):
         )
 
 
+from functools import wraps
+
 from burr.core.typing import ActionSchema
 
+
+def type_eraser(func: Callable[..., Any]) -> Callable[..., Any]:
+    """Decorator for ``run``, ``stream_run``, and ``run_and_update`` overrides
+    that declare explicit parameters instead of ``**run_kwargs``.
+
+    Applying this decorator prevents mypy ``[override]`` errors caused by
+    narrowing the base-class signature (which uses ``**run_kwargs``).
+
+    Example usage::
+
+        from burr.core import Action, State, type_eraser
+
+        class Counter(Action):
+            @property
+            def reads(self) -> list[str]:
+                return ["counter"]
+
+            @type_eraser
+            def run(self, state: State, increment_by: int) -> dict:
+                return {"counter": state["counter"] + increment_by}
+
+            @property
+            def writes(self) -> list[str]:
+                return ["counter"]
+
+            def update(self, result: dict, state: State) -> State:
+                return state.update(**result)
+
+            @property
+            def inputs(self) -> list[str]:
+                return ["increment_by"]
+    """
+
+    if inspect.iscoroutinefunction(func):
+
+        @wraps(func)
+        async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+            return await func(*args, **kwargs)
+
+        return async_wrapper
+
+    if inspect.isasyncgenfunction(func):
+
+        @wraps(func)
+        async def async_gen_wrapper(*args: Any, **kwargs: Any) -> Any:
+            async for item in func(*args, **kwargs):
+                yield item
+
+        return async_gen_wrapper
+
+    if inspect.isgeneratorfunction(func):
+
+        @wraps(func)
+        def gen_wrapper(*args: Any, **kwargs: Any) -> Any:
+            yield from func(*args, **kwargs)
+
+        return gen_wrapper
+
+    @wraps(func)
+    def wrapper(*args: Any, **kwargs: Any) -> Any:
+        return func(*args, **kwargs)
+
+    return wrapper
+
+
 # This is here to make accessing the pydantic actions easier
 # we just attach them to action so you can call `@action.pyddantic...`
 # The IDE will like it better and thus be able to auto-complete/type-check