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.

Answer 2

It's clear from @Peter's answer that we cannot make entire block as noop . I ended up doing following for my use case.

# Module foo.py
class Foo(object):
    def __init__(self):
        print "class inited"

    def __enter__(self):
        print "entered class"
        return None

    def foo(self):
        raise Exception("Not implemented")

    def __exit__(self, type, value, traceback):
        print "exited class"
        return True

----------------------------
# Module FooTest
import foo

class FooTest(object):
    def __init__(self):
        print "class inited"

    def __enter__(self):
        print "entered class"
        return None

    def __exit__(self, type, value, traceback):
        print "exited class"
        return True

try:
    foo.Foo()
    print "It shouldn't print"
except:
    print "Expected exception"
setattr(foo, 'Foo', FooTest)
print "Patched"
with foo.Foo() as a:
    a.foo()
    print "It shouldn't print"
print 'Test passed!'