How do I write a decorator that restores the cwd?

心不动则不痛 提交于 2019-11-28 19:13:38
CharlesB

The path.py module (which you really should use if dealing with paths in python scripts) has a context manager:

subdir = d / 'subdir' #subdir is a path object, in the path.py module
with subdir:
  # here current dir is subdir

#not anymore

(credits goes to this blog post from Roberto Alsina)

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.

codeape

The given answers fail to take into account that the wrapped function may raise an exception. In that case, the directory will never be restored. The code below adds exception handling to the previous answers.

as a decorator:

def preserve_cwd(function):
    @functools.wraps(function)
    def decorator(*args, **kwargs):
        cwd = os.getcwd()
        try:
            return function(*args, **kwargs)
        finally:
            os.chdir(cwd)
    return decorator

and as a context manager:

@contextlib.contextmanager
def remember_cwd():
    curdir = os.getcwd()
    try:
        yield
    finally:
        os.chdir(curdir)
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
易学教程内所有资源均来自网络或用户发布的内容,如有违反法律规定的内容欢迎反馈
该文章没有解决你所遇到的问题?点击提问,说说你的问题,让更多的人一起探讨吧!