[Python-ideas] chdir context manager

[Daniel Shahaf]

The following is a common pattern (used by, for example,
shutil.make_archive):

    save_cwd = os.getcwd()
    try:
        foo()
    finally:
        os.chdir(save_cwd)

I suggest this deserves a context manager:

    with saved_cwd():
        foo()

Initial feedback on IRC suggests shutil as where this functionality
should live (other suggestions were made, such as pathlib).  Hence,
attached patch implements this as shutil.saved_cwd, based on os.fchdir.

The patch also adds os.chdir to os.supports_dir_fd and documents the
context manager abilities of builtins.open() in its reference.

Thoughts?

Thanks,

Daniel


diff -r 74b0461346f0 Doc/library/functions.rst
--- a/Doc/library/functions.rst	Fri Jan 18 17:53:18 2013 -0800
+++ b/Doc/library/functions.rst	Sat Jan 19 09:39:27 2013 +0000
@@ -828,6 +828,9 @@ are always available.  They are listed h
    Open *file* and return a corresponding :term:`file object`.  If the file
    cannot be opened, an :exc:`OSError` is raised.
 
+   This function can be used as a :term:`context manager` that closes the
+   file when it exits.
+
    *file* is either a string or bytes object giving the pathname (absolute or
    relative to the current working directory) of the file to be opened or
    an integer file descriptor of the file to be wrapped.  (If a file descriptor
diff -r 74b0461346f0 Doc/library/os.rst
--- a/Doc/library/os.rst	Fri Jan 18 17:53:18 2013 -0800
+++ b/Doc/library/os.rst	Sat Jan 19 09:39:27 2013 +0000
@@ -1315,6 +1315,9 @@ features:
    This function can support :ref:`specifying a file descriptor <path_fd>`.  The
    descriptor must refer to an opened directory, not an open file.
 
+   See also :func:`shutil.saved_cwd` for a context manager that restores the
+   current working directory.
+
    Availability: Unix, Windows.
 
    .. versionadded:: 3.3
diff -r 74b0461346f0 Doc/library/shutil.rst
--- a/Doc/library/shutil.rst	Fri Jan 18 17:53:18 2013 -0800
+++ b/Doc/library/shutil.rst	Sat Jan 19 09:39:27 2013 +0000
@@ -36,6 +36,19 @@ copying and removal. For operations on i
 Directory and files operations
 ------------------------------
 
+.. function:: saved_cwd()
+
+   Return a :term:`context manager` that restores the current working directory
+   when it exits.  See :func:`os.chdir` for changing the current working
+   directory.
+
+   The context manager returns an open file descriptor for the saved directory.
+
+   Only available when :func:`os.chdir` supports file descriptor arguments.
+
+   .. versionadded:: 3.4
+
+
 .. function:: copyfileobj(fsrc, fdst[, length])
 
    Copy the contents of the file-like object *fsrc* to the file-like object *fdst*.
diff -r 74b0461346f0 Lib/os.py
--- a/Lib/os.py	Fri Jan 18 17:53:18 2013 -0800
+++ b/Lib/os.py	Sat Jan 19 09:39:27 2013 +0000
@@ -120,6 +120,7 @@ if _exists("_have_functions"):
 
     _set = set()
     _add("HAVE_FACCESSAT",  "access")
+    _add("HAVE_FCHDIR",     "chdir")
     _add("HAVE_FCHMODAT",   "chmod")
     _add("HAVE_FCHOWNAT",   "chown")
     _add("HAVE_FSTATAT",    "stat")
diff -r 74b0461346f0 Lib/shutil.py
--- a/Lib/shutil.py	Fri Jan 18 17:53:18 2013 -0800
+++ b/Lib/shutil.py	Sat Jan 19 09:39:27 2013 +0000
@@ -38,6 +38,7 @@ __all__ = ["copyfileobj", "copyfile", "c
            "unregister_unpack_format", "unpack_archive",
            "ignore_patterns", "chown", "which"]
            # disk_usage is added later, if available on the platform
+           # saved_cwd is added later, if available on the platform
 
 class Error(OSError):
     pass
@@ -1111,3 +1112,20 @@ def which(cmd, mode=os.F_OK | os.X_OK, p
                 if _access_check(name, mode):
                     return name
     return None
+
+# Define the chdir context manager.
+if os.chdir in os.supports_dir_fd:
+    class saved_cwd:
+        def __init__(self):
+            pass
+        def __enter__(self):
+            self.dh = os.open(os.curdir,
+                              os.O_RDONLY | getattr(os, 'O_DIRECTORY', 0))
+            return self.dh
+        def __exit__(self, exc_type, exc_value, traceback):
+            try:
+                os.chdir(self.dh)
+            finally:
+                os.close(self.dh)
+            return False
+    __all__.append('saved_cwd')
diff -r 74b0461346f0 Lib/test/test_shutil.py
--- a/Lib/test/test_shutil.py	Fri Jan 18 17:53:18 2013 -0800
+++ b/Lib/test/test_shutil.py	Sat Jan 19 09:39:27 2013 +0000
@@ -1276,6 +1276,20 @@ class TestShutil(unittest.TestCase):
         rv = shutil.copytree(src_dir, dst_dir)
         self.assertEqual(['foo'], os.listdir(rv))
 
+    def test_saved_cwd(self):
+        if hasattr(os, 'fchdir'):
+            temp_dir = self.mkdtemp()
+            orig_dir = os.getcwd()
+            with shutil.saved_cwd() as dir_fd:
+                os.chdir(temp_dir)
+                new_dir = os.getcwd()
+                self.assertIsInstance(dir_fd, int)
+            final_dir = os.getcwd()
+            self.assertEqual(orig_dir, final_dir)
+            self.assertEqual(temp_dir, new_dir)
+        else:
+            self.assertFalse(hasattr(shutil, 'saved_cwd'))
+
 
 class TestWhich(unittest.TestCase):