How to Build Custom Context Managers in Python | Towards Data Science

Programmers often work with external resources such as databases, files, connections, etc. Since these operations are so common, Python has built-in tools called context managers which allow programmers to efficiently handle external sources. The built-in context manager functions might be useful, but there are still many cases where you will need custom context managers.
Under the hood, context managers perform tasks to clean-up any used connections or resources so that they are ready for later. A common use case you might have seen is a context manager that opens a file and closes it after some operation is done:
with open('sample.txt', 'w') as file:    file.write('Hello World!')
This is the most common pattern to use context managers but there are many others too:
When you are in one of these situations, you will find there are not many built-in functions that come to your aid. Thankfully, you can build one yourself!
There are two ways to build user-defined context-managers:
class-based
function-based
Since the first way is a little more complex and some readers may not be familiar with OOP concepts in Python, we will choose an equally powerful function-based method.
In this method, there are 5 steps to implement your own context-managers:
Define a function.(optional) Write any setup code your context needs.Use the yield keyword.(optional) Write any teardown code your context needs.Add the @contextlib.contextmanager decorator.
Define a function.
(optional) Write any setup code your context needs.
Use the yield keyword.
(optional) Write any teardown code your context needs.
Add the @contextlib.contextmanager decorator.
Let’s create a simple function:
We created a very simple context manager that prints Welcome whenever we open the context and as soon as we are outside the context, prints Bye. Let's use it before I explain it further:
First of all, to convert a function to a context manager, we need to import contextlib module and add @contextlib.contextmanager decorator right before the function definition. You don't have to know what decorators are for the scope of this article but here is a link to a good source if you want to know more.
The second most important step is to add a yield keyword. The code you write before yield is called setup. In the setup, you should write all the logic that should be executed within the context like connecting to a database, opening a file, etc.
What comes after yield is always executed as soon as you are out of the context. Since Python uses indents for code blocks, writing a single unindented statement indicates to Python that we are out of the context. The code that comes after yield is also called teardown code. In the teardown, you do operations like closing connections, clearing up resources, and so on.
To drive the point home, let’s implement a commonly-used timer:
Python starts the timer whenever the interprets sees with timer(): and ends it (executes yield and what comes after) as soon as a new unindented statement is written (print('Done!'))
Sometimes, we will need context managers that accept arguments and do operations on them in the setup section. For the first example, let’s implement the built-in open function's context-manager, but only with read-only functionality:
More often than not, when context managers accept arguments they yield the results of the setup code. When a context yields a result, it can be accessed as keyword like we are doing in the above read_only function.
Here is another example that connects to a postgreSQL database and returns the connection:
It is always a good practice to implement error-handling in your code. Context managers are no exception.
Imagine in our postgres_connect function, the user generates an error before the context reaches the teardown. If this ever happens, the connection created to the database using your credentials remains open and you will no longer be able to connect again. To prevent this, we will insert try-except statements that close the connection regardless of whether there is an error or not:
After the try-except statements, adding the finally block ensures that the code within finally executes no matter what happens. This feature is particularly useful for context managers.
Now, consider this example:
The above manager temporarily changes the working directory that will allow you to perform operations within the context:
with in_dir('data'):    print(os.listdir())['sample.txt']
Of course, the above code-block could have been executed without a context manager but you get the idea. If we made any errors within the context, the error would be raised but the working directory would be back to where it was.
I think by reading this article, you can add context managers to the enormous “The reasons I love Python” pile. They are elegant solutions to many common operations that involve external resources. Whenever you see one of these patterns, be sure to use a context manager or build one for yourself!