Context Management

def contextmanager(func):
    """
    Decorator to create async context manager from generator function.

    Parameters:
    - func: Callable[..., AsyncGenerator[T, None]] - Generator function yielding once

    Returns:
    Callable[..., ContextDecorator[T]] - Function returning context manager

    Usage:
    @contextmanager
    async def my_context():
        # setup code
        yield value
        # cleanup code
    """

class ContextDecorator:
    """
    Base class for context managers usable as decorators.
    
    Inheriting from this class allows context managers to be used
    as decorators on async functions.
    """
    
    def _recreate_cm(self):
        """
        Create a copy of the context manager for decorator usage.
        
        Returns:
        Self - Copy of the context manager
        """
    
    def __call__(self, func):
        """
        Use context manager as decorator.

        Parameters:
        - func: Async callable to decorate

        Returns:
        Decorated async callable
        """

class closing:
    """
    Context manager that calls aclose() on the wrapped object when exiting.
    """
    
    def __init__(self, thing):
        """
        Parameters:
        - thing: Object with aclose() method - Resource to manage
        """
    
    async def __aenter__(self):
        """
        Enter context and return the managed resource.

        Returns:
        The managed resource object
        """
    
    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """
        Exit context and close the resource.

        Parameters:
        - exc_type: Exception type or None
        - exc_val: Exception value or None  
        - exc_tb: Exception traceback or None
        """

class nullcontext:
    """
    No-op async context manager.
    """
    
    def __init__(self, enter_result=None):
        """
        Parameters:
        - enter_result: T, optional - Value to return from __aenter__
        """
    
    async def __aenter__(self):
        """
        Enter context and return enter_result.

        Returns:
        T - The enter_result value
        """
    
    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """
        Exit context (no-op).

        Parameters:
        - exc_type: Exception type or None
        - exc_val: Exception value or None
        - exc_tb: Exception traceback or None
        """

class ExitStack:
    """
    Context manager for managing multiple async context managers.
    """
    
    def __init__(self):
        """Initialize empty exit stack."""
    
    def pop_all(self):
        """
        Remove all contexts from stack and return new stack with them.

        Returns:
        ExitStack - New stack containing all contexts
        """
    
    def push(self, exit):
        """
        Add context manager or exit callback to stack.

        Parameters:
        - exit: AsyncContextManager or exit callback - Context to manage

        Returns:
        The same context manager or callback
        """
    
    def callback(self, callback, *args, **kwargs):
        """
        Add callback function to be called on exit.

        Parameters:
        - callback: Callable - Function to call on exit
        - *args: Arguments for callback
        - **kwargs: Keyword arguments for callback

        Returns:
        Callable - The callback function
        """
    
    async def enter_context(self, cm):
        """
        Enter context manager and add to stack.

        Parameters:
        - cm: AsyncContextManager[T] - Context manager to enter

        Returns:
        T - Result from context manager's __aenter__
        """
    
    async def aclose(self):
        """Close all contexts in reverse order."""
    
    async def __aenter__(self):
        """
        Enter the exit stack context.

        Returns:
        ExitStack - This exit stack instance
        """
    
    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """
        Exit all managed contexts in reverse order.

        Parameters:
        - exc_type: Exception type or None
        - exc_val: Exception value or None
        - exc_tb: Exception traceback or None

        Returns:
        bool - True if exception was handled
        """

from asyncstdlib import contextmanager
import asyncio

@contextmanager
async def database_transaction():
    """Context manager for database transactions."""
    print("BEGIN TRANSACTION")
    try:
        yield "transaction_handle"
    except Exception:
        print("ROLLBACK")
        raise
    else:
        print("COMMIT")

async def transaction_example():
    async with database_transaction() as tx:
        print(f"Using {tx}")
        # Automatic commit on success
    
    try:
        async with database_transaction() as tx:
            print(f"Using {tx}")
            raise ValueError("Something went wrong")
    except ValueError:
        pass  # Automatic rollback on exception

from asyncstdlib import contextmanager

@contextmanager
async def timing_context():
    """Context manager that measures execution time."""
    import time
    start = time.time()
    try:
        yield
    finally:
        end = time.time()
        print(f"Execution took {end - start:.2f} seconds")

# Use as context manager
async def context_usage():
    async with timing_context():
        await asyncio.sleep(1)  # "Execution took 1.00 seconds"

# Use as decorator (if inheriting from ContextDecorator)
@timing_context()
async def timed_function():
    await asyncio.sleep(0.5)  # "Execution took 0.50 seconds"

from asyncstdlib import closing, nullcontext
import aiofiles

async def resource_example():
    # Automatic cleanup with closing
    async with closing(await aiofiles.open("file.txt")) as file:
        content = await file.read()
        # File automatically closed on exit
    
    # Conditional context manager
    use_file = True
    async with (aiofiles.open("data.txt") if use_file else nullcontext("default")) as resource:
        if use_file:
            data = await resource.read()
        else:
            data = resource  # "default"

from asyncstdlib import ExitStack
import aiofiles
import aiohttp

async def multi_context_example():
    async with ExitStack() as stack:
        # Add multiple context managers
        file1 = await stack.enter_context(aiofiles.open("input.txt"))
        file2 = await stack.enter_context(aiofiles.open("output.txt", "w"))
        session = await stack.enter_context(aiohttp.ClientSession())
        
        # Add cleanup callback
        stack.callback(lambda: print("All resources cleaned up"))
        
        # Use all resources
        data = await file1.read()
        async with session.get("https://api.example.com/data") as resp:
            api_data = await resp.text()
        
        await file2.write(f"Combined: {data} + {api_data}")
        
        # All contexts automatically exited in reverse order
        # Callback executed last

async def context_transfer():
    stack1 = ExitStack()
    
    async with stack1:
        file1 = await stack1.enter_context(aiofiles.open("temp.txt", "w"))
        await file1.write("test data")
        
        # Transfer contexts to another stack
        stack2 = stack1.pop_all()
    
    # file1 is still open, managed by stack2
    async with stack2:
        # file1 closed here
        pass

@contextmanager
async def error_handling_context():
    """Context manager with custom error handling."""
    print("Setting up resource")
    try:
        yield "resource"
    except ValueError as e:
        print(f"Handling ValueError: {e}")
        # Return True to suppress the exception
        return True
    except Exception as e:
        print(f"Unexpected error: {e}")
        # Re-raise unexpected exceptions
        raise
    finally:
        print("Cleaning up resource")

async def error_example():
    # ValueError is suppressed
    async with error_handling_context() as resource:
        raise ValueError("This will be handled")
    
    print("Execution continues")  # This runs because exception was suppressed