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)

• Anand S Kumar: if the zipFile.ZipFile becomes no-op , then how would the testing work?
• Jatin Kumar: Sorry Anand, I modified my code a bit. But the aim is to log that TestClass() was called with some arguments and then it should exit the with block immediately.
• James Mills: I'm not sure I understand your question anymore.

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.

• Jatin Kumar: Thanks for posting the original code. It's quite clear that we certainly cannot make entire block as no-op. Though if acceptable, one can also get there, check my answer.

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!'