python-logo

Python coding tips #1: with statements

on Jan 26, 16 • by Trevor Reid • with No Comments

A cool thing that is a bit underused in Python applications base is the usage of the with statement...

Home » General Coding » Python coding tips #1: with statements

A cool thing that is a bit underused in Python applications base is the usage of the with statement.

The official doc for it can be found here, it reads:

“The with statement is used to wrap the execution of a block with methods defined by a context manager (see section With Statement Context Managers). This allows common tryexceptfinally usage patterns to be encapsulated for convenient reuse.”

As described, the with statement replaces commonly used try/except/finally for nicer looking code. The most common example of this built into Python is for file handlers.

Most common usage: File handler

When manipulating files in Python, your code breaks down into the following steps:

1. open() a file
2. do stuff with file
3. close() open files

For example:

def read_log_file(self, test_file):
    f = open(test_file, "r")
    lines = f.readlines()
    f.close()
 
    return lines

Pretty simple, but it can be simplified to:

def read_log_file(self, test_file):
    with open(test_file, "r") as f:
        lines = f.readlines()

    return lines

Not much difference in this case, but when using with, the call to f.close() is handled for you, so no matter how the with block exits, the file is guaranteed to be closed. The above block is equivalent to:

def read_log_file(self, test_file):
    try:
        f = open(test_file, "r")
        lines = f.readlines()
    except:
        raise Exception
    finally:
        f.close()
    return lines 
 

Again, the main benefit of this is you can rely on the handler to close the file for you instead of having to, or maybe forgetting to, close the file yourself.

So how does it work?

To take advantage of the with statement, a class must implement an __enter__() and an __exit__() method. When with is called, it will call the class’ __enter__() method and upon exiting the with block, the class’ __exit__() method will be called. If you set up your class with these methods, then you can take advantage of with and save on coding. For an example, lets say we have a class that copies over installer files from a build server, which is used in our installer class.

We can define a copy class like so:

import shutil
import os
import tempfile
 
class CopyInstallers(object):
 
    SERVER_INSTALLER = "vncserver.exe"
    VIEWER_INSTALLER = "vncviewer.exe"
    def __init__(self, version):
        self.version = version
        self.temp_installer_dir = tempfile.mkdtemp()
        self.installer_list = list()
    
    def __enter__(self):
        self.copy_from_server()
        return self
 
    def __exit__(self, exc_type, exc_val, exc_tb):
        shutil.rmtree(self.temp_installer_dir)
 
    def copy_from_server(self):
        build_server = r"\\BUILDSERVER\release\{0}\{1}"
        for installer in [SERVER_INSTALLER, VIEWER_INSTALLER]:
            shutil.copy(build_server.format(self.version, installer), self.temp_installer_dir)
            self.installer_list.append(os.path.join(self.temp_installer_dir, installer))  

Then in your test you can simply do:

import subprocess
from copy_installer import CopyInstallers
 
 
class Install(object):
    def install(self, product, version):
        with CopyInstallers(installer=product, version=version) as copy:
            for installer in copy.installer_list:
                subprocess.call("install {0}".format(installer), shell=True)

The install() method will create a CopyInstallers class using the with keyword to copy the .exe from the server to a temporary directory for local installation. When the with statement finishes, the temporary directory will be deleted. The __exit__() method can also catch exceptions and do more error handling if needed.

With the syntax withas , the as clause will return whatever the __enter__() method returns. So in the above example, I am creating a CopyInstallers object, and want to use a reference of that object in the with body. To achieve this, I just have __enter__() return a reference to self. You can also omit the as clause if you don’t need the reference to the statement used after the with keyword.

Side note: these don’t replace __init__() or __del__() calls. __init__() will still run before __enter__() and __del__() will run after __exit__().

Known issues

One known issue with this approach is when an error occurs during the __enter__() method. If an exception is thrown here, then __exit__() will not be called. try/except blocks will have to be added in the __enter__(), and call the __exit__() method in the except/finally block, or handle the exception in another way.

Related Posts

Leave a Reply

Your email address will not be published. Required fields are marked *

Scroll to top