Created
November 14, 2024 06:18
-
-
Save mrdmnd/2f7f514abf9fa2a52ac47a88c7a42021 to your computer and use it in GitHub Desktop.
Decorator to asynchronously shut off STDOUT and STDERR while executing code.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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