from typing import Callable

# Write a simple context manager to be used in a ‹with› block. The
# goal is to enrich stack traces with additional context, like this:

def context( *args ):
    pass

# For example:

def foo( x: int, y: int ) -> None:
    with context( "asserting equality", x, '=', y ):
        assert x == y

# Calling ‹foo( 1, 1 )› should print nothing (the assertion does not
# fail and no exceptions are thrown). Ont the other hand,
# ‹foo( 7, 8 )› should print something like this:

#     asserting equality 7 = 8
#     Traceback (most recent call last):
#       File "with.py", line 20, in <module>
#         foo( 7, 8 )
#       File "with.py", line 17, in foo
#         assert x == y
#     AssertionError

def test_main() -> None:

    def foo( x: int, y: int ) -> None:
        with context( "asserting equality", x, '=', y ):
            assert x == y

    def test() -> None:
        foo( 1, 1 )
        foo( 7, 8 )

    output = redirect_err( test )

    res = output.split( '\n' )
    assert res[0] == "asserting equality 7 = 8"
    assert res[1] == "Traceback (most recent call last):"
    assert res[2].strip().startswith( "File" )
    assert res[2].endswith( "in redirect_err" )
    assert res[3] == "    f()"
    assert res[-2] == "AssertionError"

def redirect_err( f: Callable[ [], None ] ) -> str:

    import sys
    import traceback
    from io import StringIO

    stderr = sys.stderr
    out = StringIO()
    sys.stderr = out

    try:
        f()
    except:
        traceback.print_exc()

    sys.stderr = stderr
    return out.getvalue()

if __name__ == "__main__":
    test_main()