Python - How do I write a decorator that restores the cwd?

Question

How do I write a decorator that restores the current working directory to what it was before the decorated function was called? In other words, if I use the decorator on a function that does an os.chdir(), the cwd will not be changed after the function is called.

Answer 1

def preserve_cwd(function):
   def decorator(*args, **kwargs):
      cwd = os.getcwd()
      result = function(*args, **kwargs)
      os.chdir(cwd)
      return result
   return decorator

Here's how it's used:

@preserve_cwd
def test():
  print 'was:',os.getcwd()
  os.chdir('/')
  print 'now:',os.getcwd()

>>> print os.getcwd()
/Users/dspitzer
>>> test()
was: /Users/dspitzer
now: /
>>> print os.getcwd()
/Users/dspitzer

Answer 2

The answer for a decorator has been given; it works at the function definition stage as requested.

With Python 2.5+, you also have an option to do that at the function call stage using a context manager:

from __future__ import with_statement # needed for 2.5 ≤ Python < 2.6
import contextlib, os

@contextlib.contextmanager
def remember_cwd():
    curdir= os.getcwd()
    try: yield
    finally: os.chdir(curdir)

which can be used if needed at the function call time as:

print "getcwd before:", os.getcwd()
with remember_cwd():
    walk_around_the_filesystem()
print "getcwd after:", os.getcwd()

It's a nice option to have.

EDIT: I added error handling as suggested by codeape. Since my answer has been voted up, it's fair to offer a complete answer, all other issues aside.