Now that we know the structure of implementing our own class-based context manager. Let’s walk through a context manager that manages actual files as well as explore each of the methods we saw earlier. Here is what our context manager will look like:
class WorkWithFile: def __init__(self, file, mode): self.file = file self.mode = mode def __enter__(self): self.opened_file = open(self.file, self.mode) return self.opened_file def __exit__(self, *exc): self.opened_file.close()
We have written a class-based context manager called WorkWithFile
! Let’s break down each method and what happens inside of it.
The
__init__
method:This method is standard across most classes, even ones that are not context managers themselves. In this case, we have three parameters:
self
: This is standard for any class we work with and allows us to work with methods and properties we assign to an instance of a class.file
: Since we are working with files, we need to be able to take in a file argument when we call the class with awith
statement.mode
: Lastly, we need to provide the file a mode. This allows us to manage what our context manager will actually be doing, such as reading, writing, or both!
Both
file
andmode
arguments allow us to accomplish the following syntax:with WorkWithFile('file.txt', 'r')The
__enter__
method:This is where we deal with opening the file we want to work on. Since any new instance of our context manager will have a
file
andmode
property, we can pass them into theopen()
function to open a specific file with a specific mode. Then, we save it as a variable calledself.opened_file,
and return it.By returning
self.opened_file
, the file will be passed into the variable we define when we call it with thewith
statement. So for example:with WorkWithFile('file.txt', 'r') as fileWill assign the open file
'file.txt'
to the variable calledfile
that follows theas
clause and thus allowing us to use it in thewith
statement code block (which we will look at shortly).The
__exit__
method:Lastly, but one of the most important steps, we have to close the file we work on. Here we are still taking in a
*exc
argument, but we won’t touch on that until the next exercise. For now, this method is solely responsible for closing the resource we opened in__enter__
.
Now that we created our context manager, we can now use it in a with
statement like so:
with WorkWithFile("file.txt", "r") as file: print(file.read())
Open this for a step by step breakdown of the above code!
Okay, here we go:
- The
with
statement executes and a context manager object is created from the class provided, with the values of the arguments required in the__init__
method. - The
__enter__
method is run. - The
__enter__
method opens the file and returns a handle (reference) to the target variablefile
in thewith
statement. This is why in the method wereturn self.opened_file
. - Then the code block in the
with
statement executes; printing the contents of thefile
to the console. - Once the code in the
with
statement is completed, The__exit__
method is called. - The
__exit__
method closes the file.
Phew, these context managers are a handful! Take some time to let it all sink in, then let’s create our own context manager that works on poem files!
Instructions
Let’s build our poem context manager from earlier again! This time we will allow it to work on files. By the end of these exercises, we will have a custom context manager that has written to a file!
Create a class called PoemFiles
and give it a __init__
method that defines a self
, poem_file
, and mode
parameter.
Inside the method, print 'Starting up a poem context manager'
Next, let’s built the properties of the class via __init__
. Remember this is so we can pass a file name and a mode when we call the context manager with the with
statement.
Inside of the __init_
method and under the print statement, assign two properties to the class:
file
that is equal to thepoem_file
parametermode
that is equal to themode
parameter
Next, let’s work on the __enter__
method to set up what happens when we want to start working on a file.
Create an __enter__
method. Have the method print 'Opening poem file'
.
In the __enter__
method, we will need to open the file we want to work on and return it! This way, it will be assigned to the variable we declare when we work with the with
statement.
Inside __enter__
method give the class a new property called opened_poem_file
and assign it to a call of the open()
function that takes two arguments:
self.file
: our classes file propertyself.mode
: our classes mode property
Lastly, return the opened_poem_file
property!
Lastly, we need to create an __exit__
method.
Write a __exit__
method that defines a self
parameter and a *exc
parameter. Make the method print 'Closing poem file'
.
We need to make sure we close our file in the __exit__
method so we properly manage our resources.
In the __exit__
method, under the print statement, call the .close()
built-in function on the opened_poem_file
property of the class.
Uncomment the with
statement and run the code. Check out poem.txt
to see if the poem has been added!
Note: in reality, we wouldn’t have to create a context manager that opens a file because there’s already an open()
built-in function that you can run with a with
statement that will open and close a file. However, open()
has its limitations, and knowing this base structure will allow us to create our own custom and more advanced context managers that can do much more than open()
!