Building A Custom Context Manager In Python: A Closer Look

Key Takeaways

there's only a few managers i actually like one of them is the seagull manager fly in dump on everyone make a big mess and then leave i'm still trying to perfect my sequel management moves but that's not what this video is about this video is about another manager i like which is the context manager in python i'll show you how to use it and cover several ways to create your own and then i'll talk about when you should use context managers and when you shouldn't before we start i have something for you it's a free guide to help you make better design decisions you can get it at ioncode.com design guide it's pdf file pretty short to the point explaining the steps that i take whenever i design a new piece of software and hopefully it's also helpful to you so aryancodes.com design guide and the link is also in the description of this video now let's dive in the example that i'm going to start with today is pretty simple it's just a basic main function that makes a connection with a sql lite database that's stored in a file then we retrieve a cursor from that database connection we execute a query which is selecting everything from a table containing blogs and then we're going to print out the result of that query and then finally we're closing the connection so let me run this to show you the result of this so we're getting back an array containing two blocks with a title some content and so on as you can see i'm using a try finally block here and that's because i want to make sure that the connection is always closed let's say i print something here like logging.info closing connection like so then i run this again now obviously there is no error here so the connection is properly closed but if i do something else here like change the name of the table so the table block doesn't exist and then run this again then we're going to get an error and we get the trace back here but you see we're also still closing the connection that's because we have this finally clause here the problem is that here obviously we have a very simple query but it's going to be kind of annoying if we have to write everything we do with the sqlite database between these try except finally blocks and context managers are actually really helpful in that regard what are context managers context managers allow you to control setup and tear down of any sort of resource in a way that minimizes the chance of you forgetting to do the tear now especially in the presence of exceptions if you've ever opened the file in python you've probably used the with statement to manage the file this is a context manager and it's a good example of managing resources because the with statement automatically closes the file for you when you don't need it anymore so let's look a bit closer at what's actually happening in a width statement this is the syntax of the with statement with expression as variable do something behind the scenes actually a lot is happening there's an enter part and an exit part and whenever there's an exception or the block of code has completed the exit part is properly called so you don't have to worry about that in the example we saw in the beginning there was a sequel light connection that connection can actually itself be used as a context manager like so the advantage here is that the connection will automatically commit or roll back transactions however the object does not automatically close the connection so you still need to close that manually afterwards and then you'd have to make sure that that happens in all cases especially if an exception is erased the nice thing about python is that we can create our own context managers so let's create one for the sqlite example that automatically closes the connection for us what i'm going to show you is two ways of creating a context manager the hard way and the easy way now the hard way is that you create a context manager class that handles closing the connection for you and this is how you do that so let's define a class here called sqlite that's going to form the context manager for our sqlite connection and this class is going to have an initializer and that initializer will get a file name which is a string and we're going to store that as an instance variable and then we're also going to create a connection here self.connection equals sqlite 3 dot connect with file name there we go so that creates our sqlite connection now this doesn't really establish the connection that actually happens when you request a cursor so in the second part of this context manager we're going to write the enter method and in the enter method we're defining what should happen when we enter this particular context and in this case that's going to be returning a cursor to our connection so we're going to return self.connection.cursor let's add a log here like so so that we can see that enter is actually being called and we can also define what should happen when we release this resource and that's basically these things so we added a logging.info to call exit so we can see that it's actually happening we're committing and changes and then we close the connection so that's the exit method and now we have this it becomes basically really easy to use so we won't need this try except finally blocks here anymore in the main function but we can simply use the with statement now so now i just write with sqlite and then we're going to pass the file name let's also use keyword argument here as cursor so that's going to give us our cursor the try block that's no longer needed this is no longer needed then we can call cursor.execute select all from blogs and the whole finally part we can now delete because that's handled by the context manager and now let's run this again and now you see we have calling enter then it retrieves the blocks and then it's calls exits and let's say i make the same mistake again and i write the wrong table name here so let's run this again and then you see we get of course a sqlite's error that there is no such tape but if you go up you see that we're still calling exit and cleaning up the resource which is exactly what we want so because we have now the context manager it's way easier to use sqlite stuff anywhere in our code because we just use the with statement and we're sure that the connection is actually nicely closed when we're done now we can also set this up slightly differently so instead of creating this class here with the donder methods there's also a way to do the same thing but with decorator so let me delete this class for now and what i'm going to do is use context lib and we're going to use the context manager decorator to do something very similar so i'm going to write here context manager which is the decorator that we use and then i have a function let's call that open database that gets a file name just like our class context manager that we had before and now i can write this you see my github co-pilot actually already creates this code for me so we have here a connection so we're connecting with the database then we get the cursor which should actually be inside the try block because otherwise why have a try block right so we try to get the cursor and then we're going to yield that cursor and because we're using the yield operation here open db actually acts like a generator so we can open a database and then use the cursor and then whenever we get a kind of exception we're going back into this function here and then it's still going to catch the exception here and clean up the mess that we made and then we have the same thing as before so in the finally part we're closing the connection and we can even add here just like we did before a commit function call so that any changes are actually also being committed and currently this is locking the error but you could actually also remove this and just set it up like this very minimal setup actually so this is our new way or maybe even a shorter way of setting up a context manager and in the main part well we're basically using it in almost the same way but now of course we're not using the class but we're using the open db function call and that's going to yield a cursor so that works in exactly the same way as we had before so let's save this and let's now run this code and see what happens so here you see we have the error but it's still going to close the connection maybe i should add a log here closing connection like so let's run this again so now if i move up you see that we're still closing the connection here because that's what the context manager is doing for us and the syntax here is more or less the same we still do with statements until really practical let's change this back to blocks so we still know that it works and now you see that we get the blocks that we want to have and then we're also closing the connection as we expect it to so that's the decorator version of the context manager and this also works this is a pretty short way to do it i really like this syntax it's really easy to understand what's happening here as opposed to the class version which is a bit more complicated with the donder enter and exit methods context managers also support async and await syntax by the way if you want to learn more about async concurrent operations in python check out this video that i did recently about that topic and here's a simple example of an asynchronous context manager taken from the python documentation specifically for sql lite there's actually a package called aeo sqlite not sure how to pronounce that that handles both asynchronous operations as well as automatic closing of the connection so let me just quickly show you how that works and that basically means that in this particular situation you don't need to write your own context manager so that's i o sql lite and then instead of this with statement what you can also do is you write here you see main is now async right so i can uh call a weight here and then i'm calling it with connect and that should be application.db so so that's my database and what i can do then is right here cursor equals database dot execute select everything from blocks and then we're logging whatever the cursor gets also light actually also allows to use a with statement and then basically this the database is actually a context manager and then you can use the database within the with statement so before we can run this actually fetch all is also an asynchronous method so we also need to write the weight in front of this and same thing for the execute method that's something you always have to get used to when you write a synchronous code that you need to not forget to write awaits in front of things so connecting with database executing query and then fetching the results let's run this and now you see we get again our two blocks one final thing that i want to show you is that the database connection here is actually also a context manager so you could also do something like this so we write async with and then sorry i need to type correctly async with ao sqlite.connect as database and then we have these two lines that go within the with statements so within the with statement we have now the database which is the context manager and then we run again these two lines of code and then when you run this program you're going to get exactly the same result and even the cursor itself is also again a context manager so you could even write something like this so within the with statement we could write async with and then basically anything that's written here as cursor like so and then this line we need to indent of course and then basically this is what you get and then you still get the same result now i think these nested with statements are not great so i would kind of avoid that and you can probably also add it here in on a single line so i could also do it like this so db dot execute select star from blogs as cursor like so and then we don't need this and this we can de-indent so then this is what we get and then when i run this then this is still going to work again so lots of varieties and ways in which you can do this so this does mean actually that the asynchronous with it runs the various asynchronous tasks that you specify here in sequence so it's not doing that in parallel otherwise we couldn't call db.execute on the db variable but in this way you can still do it and not have all the multiple layers of indentation and python even allows you to write these things between parentheses for clarity so i could write here a parenthesis and then i could write another one here and then it can be on multiple lines like so this by the way only possible in python 3.10 but then again we get exactly the same result but it's a bit easier to see what's happening so that's asynchronous context managers pretty complicated topic but still it's nice that it's possible to do these things in python so when should you and when shouldn't you use a context manager well the main reason to use a context manager is to help allocate and most importantly clean up resources this is helpful if you want to open a file database connection a network connection if an error occurs the context manager mechanism will automatically clean up the resource for you this also means that if you don't need to clean up resources there's actually no reason to use a context manager because context managers use the with statement in python this results in that case an unnecessary code annotation level and a new scope as well because the variable is not accessible outside of the with statement so i hope you enjoyed this video if you did give it a like consider subscribing to my channel if you want to learn more about software design and development thanks for watching take care and see you soon