Suppose I have some kind of context manager (from a third-party library) that I am using like so:
with freeze_time(test_dt):
lines_of_code_1
lines_of_code_2
lines_of_code_3
But, suppose if there is no value for test_dt, the context manager should not run, but all of the remaining code should run, like so:
if test_dt:
with freeze_time(test_dt):
lines_of_code_1
lines_of_code_2
lines_of_code_3
else:
lines_of_code_1
lines_of_code_2
lines_of_code_3
Assume that lines_of_code
here is 2-3 lines of code which are exactly identical, is there a cleaner way of writing this? I'm aware that I could write something like this:
def do_thing():
lines_of_code_1
lines_of_code_2
lines_of_code_3
if test_dt:
with freeze_time(test_dt):
do_thing()
else:
do_thing()
But I'm not crazy about this formatting. Also, I don't want to have to litter this pattern all over my code.
There is one final possibility, but I'm not certain it will work: subclassing the context manager and skipping the __enter__
and __exit__
functions if the test_dt given is empty, like so:
class optional_freeze_time(object):
def __init__(self, test_dt=None):
if test_dt:
self.ctx_manager = freeze_time(test_dt)
else:
self.ctx_manager = None
def __enter__(self, *args, **kwargs):
if self.ctx_manager:
self.ctx_manager.__enter__(*args, **kwargs)
def __exit__(self, *args, **kwargs):
if self.ctx_manager:
self.ctx_manager.__exit__(*args, **kwargs)
I tested it out with a blank context manager class, and it seemed to behave correctly. However, I'm worried whether a real context manager will behave correctly if I do this (I'm not very familiar with the internals of how it works).
In general, context managers and the with statement aren’t limited to resource management. They allow you to provide and reuse common setup and teardown code. In other words, with context managers, you can perform any pair of operations that needs to be done before and after another operation or procedure, such as: Open and close; Lock and release
Context managers can be written using classes or functions (with decorators). When creating context managers using classes, user need to ensure that the class has the methods: __enter__ () and __exit__ (). The __enter__ () returns the resource that needs to be managed and the __exit__ () does not return anything but performs the cleanup operations.
Depending on which one you find more readable, you might prefer one over the other. A downside of the function-based implementation is that it requires an understanding of advanced Python topics, such as decorators and generators. The @contextmanager decorator reduces the boilerplate required to create a context manager.
The yield statement itself provides the object that will be assigned to the with target variable. This implementation and the one that uses the context management protocol are practically equivalent. Depending on which one you find more readable, you might prefer one over the other.
Newer visitors may be interested in contextlib.ExitStack:
with ExitStack() as stack:
if condition:
stack.enter_context(freeze_time(...))
lines_of_code_1
lines_of_code_2
lines_of_code_3
After this with
statement, freeze_time
is only relevant on the condition being true.
Here's an easy way to wrap around an existing context manager without even using any classes:
from contextlib import contextmanager
@contextmanager
def example_context_manager():
print('before')
yield
print('after')
@contextmanager
def optional(condition, context_manager):
if condition:
with context_manager:
yield
else:
yield
with example_context_manager():
print(1)
with optional(True, example_context_manager()):
print(2)
with optional(False, example_context_manager()):
print(3)
Output:
before
1
after
before
2
after
3
I'd probably inherit from the parent context manager and write something like this:
class BaseContextManager:
def __enter__(self):
print('using Base')
def __exit__(self, *args, **kwargs):
print('exiting Base')
class MineContextManager(BaseContextManager):
def __init__(self, input=None):
self.input = input
def __enter__(self):
if self.input:
super().__enter__()
def __exit__(self, *args, **kwargs):
if self.input:
super().__exit__()
if __name__ == '__main__':
with BaseContextManager():
print('code with base')
with MineContextManager():
print('code without base')
with MineContextManager(input=True):
print('code again with base')
This gives:
using Base
code with base
exiting Base
code without base
using Base
code again with base
exiting Base
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