Safe Directory Hopping “with” Python

Some applications need hopping around a lot between different current working directories (CWD). If you are writing such one, you might get worried that at some point control is lost and the CWD might differ from what you expected. This could certainly lead to potentially harmful consequences.

Since Python 2.5 there is a really cool feature: the with statement. There is a nice article by Fredrik Lundh [1] explaining how it works. Here I present a way to keep control over your application’s CWD making use of it.

Introduction

Suppose that our CWD is main and we want to do_something() in directory main/sub. After that we want to do_something_else() back in main. If we write

os.chdir('sub')
do_something()
os.chdir('..')
do_something_else()

everything is fine. But our entire application will crash if one of the steps fails. We can easily catch exceptions raised by do_something() but what if entering the directory fails? Writing

try:
    os.chdir('sub')
    do_something()
finally:
    os.chdir('..')

do_something_else()

certainly is a bad idea since it would jump to the parent directory of main if the os.chdir('sub') fails. A safe alternative would be to write

old_dir = os.getcwd() # returns absolute path
try:
    os.chdir('sub')
    do_something()
finally:
    os.chdir(old_dir)

do_something_else()

which is already quite a lot to type. Once we are thinking of writing a tool for it, wouldn’t i be nice if we could select sub to be created if it doesn’t exist?

There are two possible scenarios of what you need. Is it more important that sub actually exists and do_something() dosn’t make sense otherwise anyhow (as assumed in the above code) or is it more important that do_something() actually gets done no matter if the directory was there or not. Maybe you want to delete it afterwards anyhow. I’ll refer to the first situation as the “soft” and to the latter as the “strong” situation. I’ll make a suggestion for either case.

Implementation

Both situations have similar needs. We need a class object that can be used inside the with statement. I will call this class InOutSoft and InOutStrong respectively. It must at least come with the method __enter__(self) that is automatically executed at the head of the with statement and a method __exit__(self, type, value, traceback) that is executed at the bottom of the statement. The latter must ensure that, under all circumstances, the CWD after the with block is the same as it was before. No matter if it was left normally or due to an exception. We also wish to have an __init__(self, directory) method that allows us to select the directory we want to hop in.

The “Soft” Way

Having the following class

import os

class InOutSoft(object):
    def __init__(self, directory):
        self.old_dir = os.getcwd()
        self.new_dir = directory

    def __enter__(self):
        return self

    def __exit__(self, type, value, traceback):
        self.go_out()
        return isinstance(value, OSError)

    def go_in(self):
        os.chdir(self.new_dir)

    def go_out(self):
        os.chdir(self.old_dir)

we can write

with InOutSoft('sub') as d:
    d.go_in()
    do_something()

do_something_else()

to ensure that do_something_else() is always executed in main. If we can enter the directory main/sub then do_something() will be done there. If the program can’t enter main/sub, the OSError will be swallowed down. Any other exception caused inside the with block will raise through. But no matter what goes wrong inside the block, the CWD after it will always be the same as it was before. This is very similar to what the built-in

with open('hello.txt', 'w') as f:
    f.write('hello world!')

does.

The “Strong” Way

If instead we write

import sys
import os
import tempfile

class InOutStrong(object):
    def __init__(self, directory):
        self.old_dir = os.getcwd()
        self.new_dir = directory
        self.success = None

    def __enter__(self):
        try:
            try:
                # see if we can enter the directory
                os.chdir(self.new_dir)
                self.success = True
            except OSError:
                try:
                    # see if we can create it
                    os.mkdir(self.new_dir)
                    os.chdir(self.new_dir)
                    self.success = True
         except OSError:
                    # fake it!
                    self.new_dir = tempfile.mkdtemp()
                    sys.stderr.write("Couldn't enter or create directory `"
                                     + str(self.new_dir) + "'. Entering "
                                     + "temporary directory instead. All "
                                     + "data written will be lost!" + '\n')
                    self.success = False
        finally:
            return self

    def __exit__(self, type, value, traceback):
        os.chdir(self.old_dir)
        try:
            if not self.success:
                shutil.rmtree(self.new_dir)
        except:
            pass
        return isinstance(value, OSError)

then we can write

with InOutStrong('sub') as d:
    do_something()

do_something_else()

We can be sure that do_something() will never fail due to OSErrors raised by failing to enter main/sub. Note that we need no call to a go_in() method here. If it matters, we could check the value of d.success before we do_something().

To see how the two solutions work, try the following: Replace do_something() and do_something_else() by print os.getcwd() and run three cases:

• directory sub exists and we have permission to enter
• directory sub does not exist (but can be created)
• directory sub exists but we have no permission to enter

References

[1] Fredrik Lundh, Understanding Python’s “with” statement. October 2006, online http://effbot.org/zone/python-with-statement.htm.