Handle async gen and regular gen

Add overloads to `monitor_leaks` to avoid pyright ignores

Remove extra layer of callable in `monitor_leaks`

Author: masenf

reflex/monitoring.py

@@ -3,7 +3,9 @@
 import asyncio
 import contextlib
 import functools
-from collections.abc import Callable
+import inspect
+from collections.abc import AsyncGenerator, Awaitable, Callable, Generator
+from typing import TypeVar, overload
 
 from reflex.config import get_config
 
@@ -91,28 +93,69 @@ async def monitor_async():
         yield
 
 
-def monitor_leaks():
+YieldType = TypeVar("YieldType")
+SendType = TypeVar("SendType")
+ReturnType = TypeVar("ReturnType")
+
+
+@overload
+def monitor_leaks(
+    func: Callable[..., AsyncGenerator[YieldType, ReturnType]],
+) -> Callable[..., AsyncGenerator[YieldType, ReturnType]]: ...
+
+
+@overload
+def monitor_leaks(
+    func: Callable[..., Generator[YieldType, SendType, ReturnType]],
+) -> Callable[..., Generator[YieldType, SendType, ReturnType]]: ...
+
+
+@overload
+def monitor_leaks(
+    func: Callable[..., Awaitable[ReturnType]],
+) -> Callable[..., Awaitable[ReturnType]]: ...
+
+
+def monitor_leaks(func: Callable) -> Callable:
     """Framework decorator using the monitoring module's context manager.
 
+    Args:
+        func: The function to be monitored for leaks.
+
     Returns:
         Decorator function that applies PyLeak monitoring to sync/async functions.
     """
+    if inspect.isasyncgenfunction(func):
 
-    def decorator(func: Callable):
-        if asyncio.iscoroutinefunction(func):
+        @functools.wraps(func)
+        async def async_gen_wrapper(*args, **kwargs):
+            async with monitor_async():
+                async for item in func(*args, **kwargs):
+                    yield item
 
-            @functools.wraps(func)
-            async def async_wrapper(*args, **kwargs):
-                async with monitor_async():
-                    return await func(*args, **kwargs)
+        return async_gen_wrapper
 
-            return async_wrapper
+    if asyncio.iscoroutinefunction(func):
 
         @functools.wraps(func)
-        def sync_wrapper(*args, **kwargs):
+        async def async_wrapper(*args, **kwargs):
+            async with monitor_async():
+                return await func(*args, **kwargs)
+
+        return async_wrapper
+
+    if inspect.isgeneratorfunction(func):
+
+        @functools.wraps(func)
+        def gen_wrapper(*args, **kwargs):
             with monitor_sync():
-                return func(*args, **kwargs)
+                yield from func(*args, **kwargs)
+
+        return gen_wrapper
 
-        return sync_wrapper  # pyright: ignore[reportReturnType]
+    @functools.wraps(func)
+    def sync_wrapper(*args, **kwargs):
+        with monitor_sync():
+            return func(*args, **kwargs)
 
-    return decorator
+    return sync_wrapper  # pyright: ignore[reportReturnType]

reflex/state.py

@@ -1787,7 +1787,7 @@ async def _process_event(
         # Get the function to process the event.
         if is_pyleak_enabled():
             console.debug(f"Monitoring leaks for handler: {handler.fn.__qualname__}")
-            fn = functools.partial(monitor_leaks()(handler.fn), state)
+            fn = functools.partial(monitor_leaks(handler.fn), state)
         else:
             fn = functools.partial(handler.fn, state)
 
@@ -1874,7 +1874,7 @@ async def _process_event(
 
             # Handle regular event chains.
             else:
-                yield await state._as_state_update(handler, events, final=True)  # pyright: ignore[reportArgumentType]
+                yield await state._as_state_update(handler, events, final=True)
 
         # If an error occurs, throw a window alert.
         except Exception as ex: