Quartz 4

23. Context Managers

11 min read

Context Managers provide a clean and convenient interface for obtaining and releasing resources.

To implement context managers all you need to do to your class is to implement the magic methods __enter__ and __exit__. Let’s create a dummy class:

class TempFile:
    def __init__(self) -> None:
        pass
 
    def __enter__(self):
        pass
 
    def __exit__(self):
        pass
  • __enter__: is what’s going to run before we return any kind of resource with our context manager.
  • __exit__: is what’s going to run when we move out of the context of the context manager.

Let’s use a commonly used context manager:

with open("text.txt", "w") as f:
    f.write("test!!!")

The open() function is a context manager and the way you can tell it’s a context manager is because we’re using the with statement. f is whatever is returned from the __enter__ statement. Then we’ll use it within the context of the context manager, meaning that anthing under the indentation of the context manager.

Let’s demistify this process with the TempFile class. With this class we want a context manager that can give me a random file on disk that I can do things with it and then, when I leave the context manager, it just cleans it up.

Let’s write the class:

from pathlib import Path
 
class TempFile: 
    def __init__(self, filename=None) -> None:
        if not filename:
            from random import sample
            from string import ascii_letters
            filename = "".join(sample(ascii_letters, 15))
        self.file = Path(filename)
 
    def __enter__(self):
        self.file.parent.mkdir(parents=True, exist_ok=True)
        if self.file.exists():
            self.file.unlink()
        self.file.touch()
        return self.file.open("w")
 
    def __exit__(self, *args):
        self.file.unlink()
  • a good optional argument here would be to specify a particular file name and we’ll do it in the __init__ method. It takes the file name and, if it isn’t provided, then it generates a new one using a combination of random and string modules. Then we’re going to convert whatever file name we currently have to a Path object under self.file.
  • __enter__ is going to run before we give control to the user. We want to do some checks against the file name to see if it exists and, if it does exist, just destroy it and then create a new one and then finally return an open file handler as the end result of this enter method. So, when we say with TempFile() as something, then something will be the open file handler.
  • __exit__ takes a few more arguments *args, but we’ll get into it later on. Finally we want to delete the file when we are out of the context manager.

Let’s use this context manager:

with TempFile() as tf:
    tf.write("This is a text!\nThis file will be gone soon!")
    tf.flush()
    import time
    time.sleep(5)

Again, after writing with TempFile(), the __init__ is going to run and it creates a new file and it stores it into self.file as a Path object. Then the other methods are explained above.

Let’s add some print statements so we can tell when each thing is happening:

from pathlib import Path
 
class TempFile: 
    def __init__(self, filename=None) -> None:
        print("Entering init!")
        if not filename:
            from random import sample
            from string import ascii_letters
            filename = "".join(sample(ascii_letters, 15))
        self.file = Path(filename)
 
    def __enter__(self):
        print("Entering enter...")
        self.file.parent.mkdir(parents=True, exist_ok=True)
        if self.file.exists():
            self.file.unlink()
        self.file.touch()
        return self.file.open("w")
 
    def __exit__(self, *args):
        print("Entering exit!")
        self.file.unlink()

Then:

with TempFile() as tf:
    tf.write("This is a text!\nThis file will be gone soon!")
    tf.flush()
    import time
    time.sleep(5)

Output:

Entering init!
Entering enter...
Entering exit!

Let’s give more details on the __exit__ method and for this we’ll define a new context manager:

class Expect:
    def __init__(self, *exc_types, message=None) -> None:
        self.exc_types = exc_types
        self.message = message
 
    def __enter__(self):
        return self
    
    def __exit__(self, exc_type, exc_message, exc_tb):
        if exc_type in self.exc_types:
            if self.message and self.message in exc_message.args:
                return True
            elif self.message:
                return
            return True
        raise Exception(f"Expected one of {self.exc_types} exceptions.")
  • __init__: we want to handle multiple kinds of expected exceptions that should happen within the context manager (*exc_types) and also verify that a certain exception message happened (message).

Now, let’s describe extensively the __exit__ method.

There are really three specific arguments that we would expect if something goes wrong within a context manager. All of these are related to exceptions. We’re going to have the “exc_type, the exc_message, and the exc_tb`.

Now, the exception type is the class of the exception that’s being raised (for example, we have things like a ZeroDivisionError, a KeyboardInterrupt, or a general Exception). However, for this context manager, we want to check if the exception that’s raised is one of the exception types defined when the user enters the context manager. We can do this simply by saying: if exception_type in self.exception_types:

But there’s an additional element here: the exception message. This provides more details about what exception occurred. For example, if someone raises an exception with a message like “Something went wrong,” the exception message is this form of the exception. It’s an object in itself.

Let’s take a look at what that looks like so we can figure out what to do with this exception message. If you say exc_message = Exception("Something went wrong"), now we have an exception message object. But how do we get the value of the message out of this exception object? Let’s look at its attributes: args, and traceback. If we hiy exc_message.args we got ('Something went wrong',)

Now we see that args is a tuple, and it contains the message we’re looking for as the first argument. What we’re really interested in is the exception message’s args. So, let’s do a check here: if self.message in exception_message.args:, then we can continue. But we also need to add an extra check to this if statement because if someone doesn’t pass in a particular message, we don’t want it to fail as well. So we can combine this check like this: if self.message and self.message in exception_message.args:, then we can continue on. The way we continue is by simply returning True. If this check doesn’t pass, meaning we have a message but it wasn’t in the args, all we need to do is return nothing.

Finally, if you don’t want to do any kind of message checking, then you can simply return True at this point because the exception that was caught is one of the expected exception types.

Now, the last thing we need to do is handle the case where no exceptions are raised, or an incorrect exception is raised. Before we try to run this, let’s explain how exception handling happens within a __exit__ method. So, if an exception is raised, the details are passed into these arguments right here. Now, what you can do when you’re developing an __exit__ method is decide which exceptions you want to handle and how you want to handle them, and which conditions you want to pass or not. If it’s something that is acceptable, you can tell Python that you’re good with it by simply returning True. However, if you’re not good with it, you can return nothing at all or raise your own exceptions.

But if you’ve returned True from an __exit__ method when an exception is raised within your context manager, Python considers that “good to go.” You’ve handled it, you’ve done your checks, and you’ve accepted that everything is in place. Optimally, you’ve done any kind of cleanup you needed to do with your resources, so it can move on.

Alright, we now have our new Expect class. Let’s give ourselves a little bit of room and try a couple of examples of this:

with Expect(ZeroDivisionError):
    5 / 0

We see that we didn’t get a ZeroDivisionError.

Now, if we were to do anything else that either would or wouldn’t raise an exception, we’ll see what happens. So, let’s try to convert the string ‘a’ into an integer:

with Expect(ZeroDivisionError):
    int("a")

Output:

---------------------------------------------------------------------------

ValueError                                Traceback (most recent call last)

Cell In[33], line 2
	  1 with Expect(ZeroDivisionError):
----> 2     int("a")


ValueError: invalid literal for int() with base 10: 'a'


During handling of the above exception, another exception occurred:


Exception                                 Traceback (most recent call last)

Cell In[33], line 1
----> 1 with Expect(ZeroDivisionError):
	  2     int("a")


Cell In[31], line 16, in Expect.__exit__(self, exc_type, exc_message, exc_tb)
	 14         return
	 15     return True
---> 16 raise Exception(f"Expected one of {self.exc_types} exceptions.")


Exception: Expected one of (<class 'ZeroDivisionError'>,) exceptions.

We see that it actually raised a ValueError. Our context manager raised its own exception saying it expected one of these classes, which is a ZeroDivisionError. We expected a ZeroDivisionError, but we didn’t get one, so we raised an exception because of that. Now, all we’d have to do to handle that ValueError is to also put that in the list, and now our context manager handles both:

with Expect(ZeroDivisionError, ValueError):
    int("a")

And if we really wanted to, we could specify a particular kind of message by stating message = something. So, if we just wanted to copy this ValueError message up here, we should be able to do just that:

with Expect(ZeroDivisionError, ValueError, message="invalid literal for int() with base 10: 'a'"):
    int("b")

Output:

---------------------------------------------------------------------------

ValueError                                Traceback (most recent call last)

Cell In[39], line 2
	  1 with Expect(ZeroDivisionError, ValueError, message="invalid literal for int() with base 10: 'a'"):
----> 2     int("b")


ValueError: invalid literal for int() with base 10: 'b'

Now, if we were to modify this in any way, say change that ‘a’ to a ‘b’, we now get an error raised. So, we can be very exact with the kinds of exceptions we want to handle with this Expect class, even down to the exact message that needs to be raised by an exception.

Okay, so before we go, I want to show you one final quick example of how to make context managers easily in Python, and that’s by using the contextlib module built into the Python standard library:

from contextlib import contextmanager
 
@contextmanager
def expensive_resource():
    resource = ["this is a complicated method"]
    try:
        yield resource
    except KeyboardInterrupt:
        print("I handled this correctly")
    finally:
        resource[0] = "cleaned up"

What I’ve done is import the contextmanager decorator from contextlib. Now, what this lets us do is define a generator, meaning it uses the yield expression.

Now, how this works is that the contextmanager decorator will make a context manager from a generator. So, everything that happens before your yield can be considered to be a combination of your __init__ and __enter__ methods. So, we are yielding this resource here, which is simply just a list that says, “This is a complicated method,” and that’s what’s going to be provided to the user. So, they can say with expensive_resource as resource, and they’ll get back this right here. Now, everything that happens within the context manager is up to the user, but when they exit the context of it, it’s going to continue on from the yield.

Now, we can do things like handle exceptions as part of this. So, say that a KeyboardInterrupt was raised at some point. We can decide to handle that. And then finally, if we need to, we can do cleanup after the yield as well.

So, here with just a few lines, we’ve mimicked a lot of the functionality that we’ve done with our previous context managers. We’ve procured some kind of resource, we’ve given it to the user, we’re handling exceptions, and then finally, we’re doing a cleanup as part of that.

So, now we have this defined, let’s go ahead and use it:

with expensive_resource() as res:
    print(res)

Output:

['this is a complicated method']

We have with expensive_resource as res, and print this out. Then exit the context manager. We see we printed out that list saying, “This is a complicated method.”.

Now, as part of the cleanup, it should have modified res. We see afterwards res is equal to cleaned up:

>>> res
['cleaned up']

Now, if we wanted to make sure that our error handling is working as well, we can raise ourselves a KeyboardInterrupt:

with expensive_resource() as res:
    print(res)
    raise KeyboardInterrupt

Output:

['this is a complicated method']
I handled this correctly

And we see that our print happened and then the KeyboardInterrupt happened, which we handled as part of our except. It said, “I handled this correctly,” and then we did our cleanup as well.

So, that’s a quick way you can write a context manager. Since I typically write a number of classes as part of my normal programming routine, I prefer to use the __enter__ and __exit__ methods manually instead of using the contextlib method, as I like to provide additional functionality within classes that people can use outside of a context manager as well. It’s just a little bit of functionality you can add on top of your existing code to give it that extra polish.

Graph View