Decorator to asynchronously shut off STDOUT and STDERR while executing code.

silence.py

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