I'd like to write a function similar to open
. I'd like to be able to call it with with
, but also without with
.
When I use contextlib.contextmanager
, it makes my function work fine with with
:
@contextmanager
def versioned(file_path, mode):
version = calculate_version(file_path, mode)
versioned_file = open(file_path, mode)
yield versioned_file
versioned_file.close()
So, I use it like this:
with versioned('file.txt', 'r') as versioned_file:
versioned_file.write(...)
How do I use it without with
:
versioned_file = versioned('file.txt', 'r')
versioned_file.write(...)
versioned_file.close()
It complains:
AttributeError: 'GeneratorContextManager' object has no attribute 'write'
The problem is that contextmanager
only provides exactly that; a context manager to be used in the with
statement. Calling the function does not return the file object, but a special context generator which provides the __enter__
and __exit__
functions. If you want both the with
statement and “normal” assignments to work, then you will have to have some object as the return value from your function that is fully usable and also provides the context functions.
You can do this pretty easily by creating your own type, and manually providing the context functions:
class MyOpener:
def __init__ (self, filename):
print('Opening {}'.format(filename))
def close (self):
print('Closing file.')
def write (self, text):
print('Writing "{}"'.format(text))
def __enter__ (self):
return self
def __exit__ (self, exc_type, exc_value, traceback):
self.close()
>>> f = MyOpener('file')
Opening file
>>> f.write('foo')
Writing "foo"
>>> f.close()
Closing file.
>>> with MyOpener('file') as f:
f.write('foo')
Opening file
Writing "foo"
Closing file.
You have this:
@contextmanager
def versioned(file_path, mode):
# some setup code
yield versioned_file
# some teardown code
Your basic problem of course is that what you yield
from the context manager comes out of the with
statement via as
, but is not the object returned by your function. You want a function that returns something that behaves like the object open()
returns. That is to say, a context manager object that yields itself.
Whether you can do that depends what you can do with the type of versioned_file
. If you can't change it then you're basically out of luck. If you can change it then you need to implement the __enter__
and __exit__
functions as specified in PEP 343.
In your example code, though, it already has it, and your teardown code is the same as what it does itself on context exit already. So don't bother with contextlib at all, just return the result of open()
.
For other examples where you do need __enter__
and __exit__
, if you like the contextlib style (and who doesn't?) you can bridge the two things. Write a function context
that's decorated with @contextmanager
and yields self
. Then implement:
def __enter__(self):
self.context = context() # if context() is a method use a different name!
return self.context.__enter__()
def __exit__(self, *args):
return self.context.__exit__(*args)
It's basically up to you whether you find this better or worse than separating out the setup code into __enter__
and the teardown code into __exit__
. I generally find it better.
If you love us? You can donate to us via Paypal or buy me a coffee so we can maintain and grow! Thank you!
Donate Us With