Monkey patch python with statement

Question

I am using py.test for my python unit testing. Consider following code:

def mytest():
    "Test method"
    print "Before with statement"
    with TestClass('file.zip', 'r') as test_obj:
        print "This shouldn't print after patching."
        # some operation on object.
    print "After with statement."

Is it possible to monkeypatch TestClass class so that code in with block becomes a noop?

For example, the output after patching should be:

Before with statement
After with statement

I know I can patch mytest function itself, but this is in an attempt to have better test coverage.

I have tried, something on following lines but couldn't get it working.

class MockTestClass(object):
    def __init__(self, *args):
        print "__init__ called."

    def __enter__(self):
        print "__enter__ called."
        raise TestException("Yeah done with it! Get lost now.")

    def __exit__(self, type, value, traceback):
        print "__exit__ called."

module_name.setattr('TestClass',  MockTestClass)

Answer 1

I think what you're trying to do is disallowed by the Python language specification.

As you can see in PEP-343 the definition of the "with" statement will not allow any attempt to exit the context early:

mgr = (EXPR)
exit = type(mgr).__exit__  # Not calling it yet
value = type(mgr).__enter__(mgr)
exc = True
try:
    try:
        VAR = value  # Only if "as VAR" is present
        BLOCK
    except:
        # The exceptional case is handled here
        exc = False
        if not exit(mgr, *sys.exc_info()):
            raise
        # The exception is swallowed if exit() returns true
finally:
    # The normal and non-local-goto cases are handled here
    if exc:
        exit(mgr, None, None, None)

There was a proposal to change this to the sor of function you would need (PEP-377) but this has been rejected.