Python context managers
Azanul Haque
Posted on April 3, 2024
The with statement in Python is primarily used for resource management. It ensures that a resource is properly initialized and released when it's no longer needed, even if exceptions occur during the execution. It can also be used for suppressing certain exceptions or outputs.
Basically, context-managers can be leveraged to create a temporary state change and revert it to original once the intended work is done.
A few common use cases for context-managers are:
- File I/O Operations: ensuring that file is properly closed after use.
- Network Connections: ensuring the connection is properly closed after use.
- Transaction Management: managing transactions - beggining, commiting, aborting, etc.
- Resource Management: automatically releasing the lock, freeing the resource after use.
Let's take an example to explore context-managers further. Let's say that we have a function which prints a message to standard output. We want to capture this output instead of releasing it to standard output.
Imagine you're a super spy with a secret message, but you don't want anyone to see it! That's where our special code comes in. It's like a magic box that hides messages until you open it again.
-
We need a secret message to hide. To do this we simply define a function which will print a string.
def secret_message(): print("This message is secret and should be suppressed!")
-
Let's create our message box. Our message box is going to be a context manager, so it has to implement the necessary enter() and exit() methods. enter() is called at the start of the context and exit() is called at the end, duh! enter() will return an object which can be accessed with
as
keyword. exit() is called with 3 arguments: exception type, exception value and traceback, these values are going to beNone
is no exception occured in the context. In enter() method, we're storing the stdout in an instance variable_original
and setting the stdout to aStringIO()
object. In exit() method, we're resetting the stdout and storing the captured output in_original
.
from io import StringIO import sys class SecretMessageBox: def __enter__(self): self._original = sys.stdout sys.stdout = StringIO() return self def __exit__(self, exc_type, exc_value, traceback): sys.stdout, self._original = self._original, sys.stdout
-
Now, let's use our box to hide a message. We're using with statement wo invoke our context manager and
as
to capture the object returned by enter() method of the context managerSecretMessageBox
into a variablesecret_box
.
print("Start hiding messages!") with SecretMessageBox() as secret_box: secret_message() print("Stop hiding messages!\n")
-
Once we've hidden the message, we want to read the message later
hidden_message = secret_box._original.getvalue() if hidden_message: print("This is what we hid:\n", hidden_message) else: print("There are no hidden messages yet.")
-
Putting it all together
from io import StringIO import sys def secret_message(): print("This message is secret and should be suppressed!") class SecretMessageBox: def __enter__(self): self._original = sys.stdout sys.stdout = StringIO() return self def __exit__(self, exc_type, exc_value, traceback): sys.stdout, self._original = self._original, sys.stdout print("Start hiding messages!") with SecretMessageBox() as secret_box: secret_message() print("Stop hiding messages!\n") hidden_message = secret_box._original.getvalue() if hidden_message: print("This is what we hid:\n", hidden_message) else: print("There are no hidden messages yet.")
-
In case we face an exception but want to ignore certain exception or certain type of exceptions.
We define aRuntimeError
, raise it from secret_message() function and handle it in exit() method.
True
returned from exit() implies the exception raised has been handled else the exception is raised further down the sttack.
NOT_REALLY_AN_EXCEPTION = RuntimeError("Just kidding") def secret_message(): print("This message is secret and should be suppressed!") raise NOT_REALLY_AN_EXCEPTION ... class SecretMessageBox: ... def __exit__(self, exc_type, exc_value, traceback): sys.stdout, self._original = self._original, sys.stdout if exc_value == NOT_REALLY_AN_EXCEPTION: return True ... try: with SecretMessageBox() as secret_box: secret_message() except Exception as e: print("Exception occurred:", e) sys.exit() finally: print("Stop hiding messages!\n")
-
Final script
from io import StringIO import sys NOT_REALLY_AN_EXCEPTION = RuntimeError("Just kidding") def secret_message(): print("This message is secret and should be suppressed!") raise NOT_REALLY_AN_EXCEPTION class SecretMessageBox: def __enter__(self): self._original = sys.stdout sys.stdout = StringIO() return self def __exit__(self, exc_type, exc_value, traceback): sys.stdout, self._original = self._original, sys.stdout if exc_value == NOT_REALLY_AN_EXCEPTION: return True print("Start hiding messages!") try: with SecretMessageBox() as secret_box: secret_message() except Exception as e: print("Exception occurred:", e) sys.exit() finally: print("Stop hiding messages!\n") hidden_message = secret_box._original.getvalue() if hidden_message: print("This is what we hid:\n", hidden_message) else: print("There are no hidden messages yet.")
Also checkout contextlib, it provides some common utilities involving the with statement.
Posted on April 3, 2024
Join Our Newsletter. No Spam, Only the good stuff.
Sign up to receive the latest update from our blog.