In this question, I defined a context manager that contains a context manager. What is the easiest correct way to accomplish this nesting? I ended up calling self.temporary_file.__enter__() in self.__enter__(). However, in self.__exit__, I am pretty sure I have to call self.temporary_file.__exit__(type_, value, traceback) in a finally block in case an exception is raised. Should I be setting the type_, value, and traceback parameters if something goes wrong in self.__exit__? I checked contextlib, but couldn't find any utilities to help with this.
Original code from question:
import itertools as it
import tempfile
class WriteOnChangeFile:
def __init__(self, filename):
self.filename = filename
def __enter__(self):
self.temporary_file = tempfile.TemporaryFile('r+')
self.f = self.temporary_file.__enter__()
return self.f
def __exit__(self, type_, value, traceback):
try:
try:
with open(self.filename, 'r') as real_f:
self.f.seek(0)
overwrite = any(
l != real_l
for l, real_l in it.zip_longest(self.f, real_f))
except IOError:
overwrite = True
if overwrite:
with open(self.filename, 'w') as real_f:
self.f.seek(0)
for l in self.f:
real_f.write(l)
finally:
self.temporary_file.__exit__(type_, value, traceback)
The easy way to create context managers is with contextlib.contextmanager. Something like this:
#contextlib.contextmanager
def write_on_change_file(filename):
with tempfile.TemporaryFile('r+') as temporary_file:
yield temporary_file
try:
... some saving logic that you had in __exit__ ...
Then use with write_on_change_file(...) as f:.
The body of the with statement will be executed “instead of” the yield. Wrap the yield itself in a try block if you want to catch any exceptions that happen in the body.
The temporary file will always be properly closed (when its with block ends).
contextlib.contextmanager works great for functions, but when I need a classes as context manager, I'm using the following util:
class ContextManager(metaclass=abc.ABCMeta):
"""Class which can be used as `contextmanager`."""
def __init__(self):
self.__cm = None
#abc.abstractmethod
#contextlib.contextmanager
def contextmanager(self):
raise NotImplementedError('Abstract method')
def __enter__(self):
self.__cm = self.contextmanager()
return self.__cm.__enter__()
def __exit__(self, exc_type, exc_value, traceback):
return self.__cm.__exit__(exc_type, exc_value, traceback)
This allow to declare contextmanager classes with the generator syntax from #contextlib.contextmanager. It makes it much more natural to nest contextmanager, without having to manually call __enter__ and __exit__. Example:
class MyClass(ContextManager):
def __init__(self, filename):
self._filename = filename
#contextlib.contextmanager
def contextmanager(self):
with tempfile.TemporaryFile() as temp_file:
yield temp_file
... # Post-processing you previously had in __exit__
with MyClass('filename') as x:
print(x)
I wish this was in the standard library...
Related
Hi guys I have a quick question about python context manager.
How do I create a context manager where it will release all memory from an array or list? Is it even possible to do something like that?
Say I have the array
b = [1, 2, 3, 4]
We can certainly use clear or del to clear/delete the array but it will still be there in memory. And then we can use gc.collect() to clear it from memory.
class FileManager():
def __init__(self, filename, mode):
self.filename = filename
self.mode = mode
self.file = None
def __enter__(self):
self.file = open(self.filename, self.mode)
return self.file
def __exit__(self, exc_type, exc_value, exc_traceback):
self.file.close()
with FileManager('test.txt', 'w') as f:
f.write('Test')
But is it possible to create a context manager class to do this? I know they are usually used for file management as I wrote the code above...just wondering how else to use context manager
If we look at python docs it states:
Most context managers are written in a way that means they can only be used effectively in a with statement once. These single use context managers must be created afresh each time they’re used - attempting to use them a second time will trigger an exception or otherwise not work correctly.
This common limitation means that it is generally advisable to create context managers directly in the header of the with statement where they are used (as shown in all of the usage examples above).
Yet, the example most commonly shared for creating context managers inside classes is:
from contextlib import ContextDecorator
import logging
logging.basicConfig(level=logging.INFO)
class track_entry_and_exit(ContextDecorator):
def __init__(self, name):
self.name = name
def __enter__(self):
logging.info('Entering: %s', self.name)
def __exit__(self, exc_type, exc, exc_tb):
logging.info('Exiting: %s', self.name)
But, when I instantiate this class, I can pass it several times to a with statement:
In [8]: test_context = track_entry_and_exit('test')
In [9]: with test_context:
...: pass
...:
INFO:root:Entering: test
INFO:root:Exiting: test
In [10]: with test_context:
...: pass
...:
INFO:root:Entering: test
INFO:root:Exiting: test
How can I create a class that fails on the second call to the with statement?
Here is a possible solution:
from functools import wraps
class MultipleCallToCM(Exception):
pass
def single_use(cls):
if not ("__enter__" in vars(cls) and "__exit__" in vars(cls)):
raise TypeError(f"{cls} is not a Context Manager.")
org_new = cls.__new__
#wraps(org_new)
def new(clss, *args, **kwargs):
instance = org_new(clss)
instance._called = False
return instance
cls.__new__ = new
org_enter = cls.__enter__
#wraps(org_enter)
def enter(self):
if self._called:
raise MultipleCallToCM("You can't call this CM twice!")
self._called = True
return org_enter(self)
cls.__enter__ = enter
return cls
#single_use
class CM:
def __enter__(self):
print("Enter to the CM")
def __exit__(self, exc_type, exc_value, exc_tb):
print("Exit from the CM")
with CM():
print("Inside.")
print("-----------------------------------")
with CM():
print("Inside.")
print("-----------------------------------")
cm = CM()
with cm:
print("Inside.")
print("-----------------------------------")
with cm:
print("Inside.")
output:
Enter to the CM
Inside.
Exit from the CM
-----------------------------------
Enter to the CM
Inside.
Exit from the CM
-----------------------------------
Enter to the CM
Inside.
Exit from the CM
-----------------------------------
Traceback (most recent call last):
File "...", line 51, in <module>
with cm:
File "...", line 24, in enter
raise MultipleCallToCM("You can't call this CM twice!")
__main__.MultipleCallToCM: You can't call this CM twice!
I used a class decorator for it so that you can apply it to other context manager classes. I dispatched the __new__ method and give every instance a flag called __called, then change the original __enter__ to my enter which checks to see if this object has used in a with-statement or not.
How robust is this? I don't know. Seems like it works, I hope it gave an idea at least.
Arguably the simplest method is mentioned two paragraphs further down in the documentation you have cited:
Context managers created using contextmanager() are also single use context managers, and will complain about the underlying generator failing to yield if an attempt is made to use them a second time
Here is the corresponding invocation for your example:
>>> from contextlib import contextmanager
>>> #contextmanager
... def track_entry_and_exit(name):
... print('Entering', name)
... yield
... print('Exiting', name)
...
>>> c = track_entry_and_exit('test')
>>> with c:
... pass
...
Entering test
Exiting test
>>> with c:
... pass
...
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
File "/usr/lib/python3.9/contextlib.py", line 115, in __enter__
del self.args, self.kwds, self.func
AttributeError: args
It's even a class although it is written as a function:
>>> type(c)
<class 'contextlib._GeneratorContextManager'>
I suggest to consider iterable class instead of context manager, like this
class Iterable:
"""Iterable that can be iterated only once."""
def __init__(self, name):
self.name = name
self.it = iter([self])
def __iter__(self):
# code to acquire resource
print('enter')
yield next(self.it)
print('exit')
# code to release resource
def __repr__(self):
return f'{self.__class__.__name__}({self.name})'
It can be iterated only one
>>> it = Iterable('iterable')
>>> for item in it:
>>> print('entered', item)
enter
entered Iterable(iterable)
exit
>>> for item in it:
>>> print('entered', item)
RuntimeError: generator raised StopIteration
Context manager can be written in this manner:
class Context:
"""Context manager that can be used only once."""
def __init__(self, name):
self.name = name
self.it = iter([self])
def __enter__(self):
print('enter')
return next(self.it)
def __exit__(self, exc_type, exc, exc_tb):
print('exit')
def __repr__(self):
return f'{self.__class__.__name__}({self.name})'
It works only once
>>> ctx = Context('context')
>>> with ctx as c:
>>> print('entered', c)
enter
entered Context(context)
exit
>>> with ctx as c:
>>> print('entered', c)
enter
StopIteration:
How can I override the built in open function such that when I call it like so...
with open(file_path, "r") as f:
contents = f.read()
The contents variable is any string I want?
EDIT: To clarify, I want to be able to just provide a string to the open function rather than a file path that will be read.
with open("foobar") as f:
contents = f.read()
print(contents)
The above should print foobar.
I am aware this is defeating the purpose of open etc but it is for testing purposes.
You can create your own file-like type and override the builtin open with your own open function.
import builtins
import contextlib
class File(object):
"""
A basic file-like object.
"""
def __init__(self, path, *args, **kwargs):
self._fobj = builtins.open(path, *args, **kwargs)
def read(self, n_bytes = -1):
data = self._fobj.read(n_bytes)
...
return data
def close(self):
self._fobj.close()
#contextlib.contextmanager
def open(path, *args, **kwargs):
fobj = File(path, *args, **kwargs)
try:
with contextlib.closing(fobj):
yield fobj
finally:
pass
You can add whatever behavior or additional logic needed to adjust the return value of read() inside File.read itself, or override the behavior entirely from a subclass of File.
Simplified for the particular case in question:
class File(str):
def read(self):
return str(self)
#contextlib.contextmanager
def open(string):
try:
yield File(string)
finally:
pass
with open('foobar') as f:
print(f.read())
Considering it is for testing purpose and you want to force the open calls to return a specific string then you can use mock_open here.
Let's say I have a module foo that has a function that reads content from a file and counts the number of lines:
# foo.py
def read_and_process_file():
with open('Pickle Rick') as f:
contents = f.read()
print('File has {n} lines'.format(n=len(contents.splitlines())))
Now in your test you can mock the open for this module and make it return any string you want:
from unittest.mock import mock_open, patch
import foo
m = mock_open(read_data='I am some random data\nthat spans over 2 lines')
with patch('foo.open', m):
foo.read_and_process_file() # prints 2
You can design your own class, as with requires an object with a defined __enter__ and __exit__ method. As that is what with does.
class my_class:
def __init__(self, *args):
print("initializing variable, got args: {}".format(args))
def __enter__(self):
print("Inside enter statement!")
return "arbitrary text"
def __exit__(self, type, value, traceback):
print("closing time, you don't have to go home")
return
with my_class(1,2,3) as my_thing:
print("inside the with block!")
print("The return from my_class __enter__ is: ", my_thing)
print("Outside with block!")
output when ran:
initializing variable, got args: (1, 2, 3)
Inside enter statement!
inside the with block!
The return from my_class __enter__ is: arbitrary text
closing time, you don't have to go home
Outside with block!
More reading here: http://effbot.org/zone/python-with-statement.htm
Any advice on what is the proper 'pythonic' way for the following function. Do I have to split it into two functions?
def readSomething(fp=None):
if fp:
return fp.read(100)
else:
with open('default.txt', 'r') as fp:
return fp.read(100)
I need something like this because the readSomething function may be called from another function that may or may not have the same file open.
For example, it may be called like this at some places:
def doSomethingWithSameFile():
with open('default.txt') as fp:
preample = fp.read(10)
more_data = readSomething(fb)
...
or like that at other places:
def init():
data = readSomething()
...
I don't think this is the right solution, but I think it's what you want.
import contextlib
def readSomething(fp=None):
with contextlib.ExitStack() as stack:
if not fp:
fp = stack.enter_context(open('default.txt'))
return fp.read(100)
I get the impression that you're going to duplicate this logic many functions like readSomething() so I'd recommend putting the ExitStack code into a decorator and wrapping the functions where you need this behavior.
You could also use a decorator. I don't use this kind of code, so the syntax below is almost certainly incomplete, but the general idea stands:
import functools
def fallback_to_default(fn):
#functools.wraps(fn)
def new_fn(fp=None, *args, **kwargs):
with contextlib.ExitStack() as stack:
if not fp:
fp = stack.enter_context(open('default.txt'))
return fn(fp, *args, **kwargs)
return new_fn
#fallback_to_default
def readSomething(fp=None):
return fp.read(100)
You could define a custom context manager that does something only if None is passed to it, but it might be overkill:
class ContextOrNone(object):
def __init__(self, obj, fn, *args, **kwargs):
if obj is not None:
self.obj = obj
self.cleanup = False
else:
self.obj = fn(*args, **kwargs)
self.cleanup = True
def __enter__(self):
return self.obj
def __exit__(self, ex_type, ex_val, traceback):
if self.cleanup:
self.obj.__exit__(ex_type, ex_val, traceback)
Or, using contextlib.contextmanager:
from contextlib import contextmanager
#contextmanager
def ContextOrNone(obj, fn, *args, **kwargs):
was_none = obj is None
try:
if was_none:
obj = fn(*args, **kwargs)
yield obj
finally:
if was_none:
obj.__exit__()
Once you have this defined, you can define readSomething as:
def readSomething(fp=None):
with ContextOrNone(fp, open, 'default.txt', 'r') as fp:
return fp.read(100)
To summarise the issue in plain language:
You may be passed an open file handle, then you want to leave it open, because it is the caller's responsibility to close that resource.
You might need to open your own resource, then it is your responsibility to close it.
This is the problem with accepting inhomogeneous argument types in Python. You're allowed to do that, but it can make your code a bit more ugly sometimes.
Context managers are just syntactic sugar for try/finally:
def readSomething(fp=None):
close_fp = False
if fp is None:
fp = open('default.txt')
close_fp = True
try:
return fp.read(100)
finally:
if close_fp:
fp.close()
To make it any more "pretty" than that, consider to change the interfaces so that you don't have to handle both reading data and doing the resource management from within the same function - refactor to make your functions have a single responsibility.
Honestly, the most Pythonic version of your code is probably just what you already have, except slightly cleaned up:
def readSomething(fp=None):
if fp:
return fp.read(100)
with open('default.txt') as fp:
return fp.read(100)
That preserves your original intent and functionality. It's clear and easy to read. Sure, it has a little repetition. If your example was simplified to the point that the repeated portion is too grotesque for you, then lift it out into its own function:
def complicatedStuff(buf, sz):
# Obviously more code will go here.
return buf.read(sz)
def readSomething(fp=None):
if fp:
return complicatedStuff(fp, 100)
with open('default.txt') as fp:
return complicatedStuff(fp, 100)
It's not Pythonic to jump through a lot of hoops just to avoid repeating yourself a little.
This doesn't use with, and closes the default or the file passed as argument,
but maybe it is still an option.
def readSomething(fp=None):
if fp is None:
fp = open('default.txt')
return (fp.read(100), fp.close)
I would like to be able to write code like this:
with obj.in_batch_mode:
obj.some_attr = "some_value"
obj.some_int = 142
...
when I want obj to wait with sending updates about itself until multiple jobs are completed. I have hooks on __setattr__ that take some time to run, and the changes can be sent together.
I do not want to use code like this, since it increases the risk of forgetting to leave batch_mode (which is what the with keyword is good for):
obj.enter_batch_mode()
obj.some_attr = "some_value"
obj.some_int = 142
...
obj.exit_batch_mode()
I have not been able to figure out how to implement this. Just typing with obj: (and simply implementing with on obj) does not read anywhere near as descriptive.
Generally, a very simple way to implement context managers is to use the contextlib module. Writing a context manager becomes as simple as writing a single yield generator. Before the yield replaces the __enter__ method, the object yielded is the return value of __enter__, and the section after the yield is the __exit__ method. Any function on your class can be a context manager, it just needs the be decorated as such. For instance, take this simple ConsoleWriter class:
from contextlib import contextmanager
from sys import stdout
from io import StringIO
from functools import partial
class ConsoleWriter:
def __init__(self, out=stdout, fmt=None):
self._out = out
self._fmt = fmt
#property
#contextmanager
def batch(self):
original_out = self._out
self._out = StringIO()
try:
yield self
except Exception as e:
# There was a problem. Ignore batch commands.
# (do not swallow the exception though)
raise
else:
# no problem
original_out.write(self._out.getvalue())
finally:
self._out = original_out
#contextmanager
def verbose(self, fmt="VERBOSE: {!r}"):
original_fmt = self._fmt
self._fmt = fmt
try:
yield self
finally:
# don't care about errors, just restore end
self._fmt = original_fmt
def __getattr__(self, attr):
"""creates function that writes capitalised attribute three times"""
return partial(self.write, attr.upper()*3)
def write(self, arg):
if self._fmt:
arg = self._fmt.format(arg)
print(arg, file=self._out)
Example usage:
writer = ConsoleWriter()
with writer.batch:
print("begin batch")
writer.a()
writer.b()
with writer.verbose():
writer.c()
print("before reentrant block")
with writer.batch:
writer.d()
print("after reentrant block")
print("end batch -- all data is now flushed")
Outputing:
begin batch
before reentrant block
after reentrant block
end batch -- all data is now flushed
AAA
BBB
VERBOSE: 'CCC'
DDD
If you are after a simple solution and do not need any nested mode-change (e.g. from STD to BATCH to VERBOSE back to BATCH back to STD)
class A(object):
STD_MODE = 'std'
BATCH_MODE = 'batch'
VERBOSE_MODE = 'verb'
def __init__(self):
self.mode = self.STD_MODE
def in_mode(self, mode):
self.mode = mode
return self
def __enter__(self):
return self
def __exit__(self, type, value, tb):
self.mode = self.STD_MODE
obj = A()
print obj.mode
with obj.in_mode(obj.BATCH_MODE) as x:
print x.mode
print obj.mode
outputs
std
batch
std
This builds on Pynchia's answer, but adds support for multiple modes and allows nesting of with statements, even with the same mode multiple times. It scales O(#nested_modes) which is basically O(1).
Just remember to use stacks for data storage related to the modes.
class A():
_batch_mode = "batch_mode"
_mode_stack = []
#property
def in_batch_mode(self):
self._mode_stack.append(self._batch_mode)
return self
def __enter__(self):
return self
def __exit__(self, type, value, tb):
self._mode_stack.pop()
if self._batch_mode not in self._mode_stack:
self.apply_edits()
and then I have these checks wherever I need them:
if self._batch_mode not in self._mode_stack:
self.apply_edits()
It is also possible to use methods for modes:
with x.in_some_mode(my_arg):
just remember to save my_arg in a stack within x, and to clear it from the stack when that mode is popped from the mode stack.
The code using this object can now be
with obj.in_batch_mode:
obj.some_property = "some_value"
and there are no problems with nesting, so we can add another with obj.in_some_mode: wherever without any hard-to-debug errors or having to check every function called to make sure the object's with-statements are never nested:
def b(obj):
with obj.in_batch_mode:
obj.some_property = "some_value"
x = A()
with x.in_batch_mode:
x.my_property = "my_value"
b(x)
Maybe something like this:
Implement helper class
class WithHelperObj(object):
def __init__(self,obj):
self.obj = obj
def __enter__(self):
self.obj.impl_enter_batch()
def __exit__(self, exc_type, exc_value, traceback):
self.obj.impl_exit_batch()
class MyObject(object):
def in_batch_mode(self):
return WithHelperObj(self)
In the class itself, implement method instead of field, to use with the with statement
def impl_enter_batch(self):
print 'In impl_enter_batch'
def impl_exit_batch(self):
print 'In impl_exit_batch'
def doing(self):
print 'doing'
Then use it:
o = MyObject()
with o.in_batch_mode():
o.doing()