Skip to content

Instantly share code, notes, and snippets.

@mrdmnd
Created November 14, 2024 06:18
Show Gist options
  • Save mrdmnd/2f7f514abf9fa2a52ac47a88c7a42021 to your computer and use it in GitHub Desktop.
Save mrdmnd/2f7f514abf9fa2a52ac47a88c7a42021 to your computer and use it in GitHub Desktop.
Decorator to asynchronously shut off STDOUT and STDERR while executing code.
import asyncio
import contextlib
import io
import sys
import threading
from functools import wraps
from typing import Any, Callable, TypeVar, Union, AsyncIterator
T = TypeVar('T')
class ThreadSafeStreamRedirect:
"""Thread-safe stream redirect that maintains separate buffers per thread"""
def __init__(self, original_stream):
self.original_stream = original_stream
self._local = threading.local()
def _get_buffer(self):
if not hasattr(self._local, 'buffer'):
self._local.buffer = io.StringIO()
return self._local.buffer
def write(self, text):
self._get_buffer().write(text)
def flush(self):
self._get_buffer().flush()
def get_output(self) -> str:
return self._get_buffer().getvalue()
@contextlib.asynccontextmanager
async def silence_streams() -> AsyncIterator[tuple[ThreadSafeStreamRedirect, ThreadSafeStreamRedirect]]:
"""
Async context manager that redirects stdout and stderr to thread-safe buffers
Returns the redirect objects so captured output can be accessed if needed
"""
stdout_redirect = ThreadSafeStreamRedirect(sys.stdout)
stderr_redirect = ThreadSafeStreamRedirect(sys.stderr)
old_stdout = sys.stdout
old_stderr = sys.stderr
try:
sys.stdout = stdout_redirect
sys.stderr = stderr_redirect
yield stdout_redirect, stderr_redirect
finally:
sys.stdout = old_stdout
sys.stderr = old_stderr
def silence(func: Callable[..., Union[Any, asyncio.Future]]) -> Callable[..., Union[Any, asyncio.Future]]:
"""
Decorator that silences stdout and stderr during function execution.
Works with both async and sync functions.
Thread-safe and maintains separate buffers per thread.
Args:
func: The function to be decorated (can be async or sync)
Returns:
The wrapped function that executes silently
"""
if asyncio.iscoroutinefunction(func):
@wraps(func)
async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
async with silence_streams() as (stdout_redirect, stderr_redirect):
return await func(*args, **kwargs)
return async_wrapper
else:
@wraps(func)
def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
async def _run():
async with silence_streams() as (stdout_redirect, stderr_redirect):
return await asyncio.get_event_loop().run_in_executor(
None, func, *args, **kwargs
)
return asyncio.get_event_loop().run_until_complete(_run())
return sync_wrapper
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment