Issue #25994: Added the close() method and the support of the context manager

serhiy-storchaka · serhiy-storchaka · commit ffe96ae10be8 · 2016-02-11T13:21:30.000+02:00
protocol for the os.scandir() iterator.
diff --git a/Doc/library/os.rst b/Doc/library/os.rst
@@ -1891,14 +1891,29 @@ features:
    :attr:`~DirEntry.path` attributes of each :class:`DirEntry` will be of
    the same type as *path*.
 
+   The :func:`scandir` iterator supports the :term:`context manager` protocol
+   and has the following method:
+
+   .. method:: scandir.close()
+
+      Close the iterator and free acquired resources.
+
+      This is called automatically when the iterator is exhausted or garbage
+      collected, or when an error happens during iterating.  However it
+      is advisable to call it explicitly or use the :keyword:`with`
+      statement.
+
+      .. versionadded:: 3.6
+
    The following example shows a simple use of :func:`scandir` to display all
    the files (excluding directories) in the given *path* that don't start with
    ``'.'``. The ``entry.is_file()`` call will generally not make an additional
    system call::
 
-      for entry in os.scandir(path):
-         if not entry.name.startswith('.') and entry.is_file():
-             print(entry.name)
+      with os.scandir(path) as it:
+          for entry in it:
+              if not entry.name.startswith('.') and entry.is_file():
+                  print(entry.name)
 
    .. note::
 
@@ -1914,6 +1929,12 @@ features:
 
    .. versionadded:: 3.5
 
+   .. versionadded:: 3.6
+      Added support for the :term:`context manager` protocol and the
+      :func:`~scandir.close()` method.  If a :func:`scandir` iterator is neither
+      exhausted nor explicitly closed a :exc:`ResourceWarning` will be emitted
+      in its destructor.
+
 
 .. class:: DirEntry
 
diff --git a/Doc/whatsnew/3.6.rst b/Doc/whatsnew/3.6.rst
@@ -104,6 +104,17 @@ directives ``%G``, ``%u`` and ``%V``.
 (Contributed by Ashley Anderson in :issue:`12006`.)
 
 
+os
+--
+
+A new :meth:`~os.scandir.close` method allows explicitly closing a
+:func:`~os.scandir` iterator.  The :func:`~os.scandir` iterator now
+supports the :term:`context manager` protocol.  If a :func:`scandir`
+iterator is neither exhausted nor explicitly closed a :exc:`ResourceWarning`
+will be emitted in its destructor.
+(Contributed by Serhiy Storchaka in :issue:`25994`.)
+
+
 pickle
 ------
 
diff --git a/Lib/os.py b/Lib/os.py
@@ -374,46 +374,47 @@ def walk(top, topdown=True, onerror=None, followlinks=False):
             onerror(error)
         return
 
-    while True:
-        try:
+    with scandir_it:
+        while True:
             try:
-                entry = next(scandir_it)
-            except StopIteration:
-                break
-        except OSError as error:
-            if onerror is not None:
-                onerror(error)
-            return
-
-        try:
-            is_dir = entry.is_dir()
-        except OSError:
-            # If is_dir() raises an OSError, consider that the entry is not
-            # a directory, same behaviour than os.path.isdir().
-            is_dir = False
-
-        if is_dir:
-            dirs.append(entry.name)
-        else:
-            nondirs.append(entry.name)
+                try:
+                    entry = next(scandir_it)
+                except StopIteration:
+                    break
+            except OSError as error:
+                if onerror is not None:
+                    onerror(error)
+                return
 
-        if not topdown and is_dir:
-            # Bottom-up: recurse into sub-directory, but exclude symlinks to
-            # directories if followlinks is False
-            if followlinks:
-                walk_into = True
+            try:
+                is_dir = entry.is_dir()
+            except OSError:
+                # If is_dir() raises an OSError, consider that the entry is not
+                # a directory, same behaviour than os.path.isdir().
+                is_dir = False
+
+            if is_dir:
+                dirs.append(entry.name)
             else:
-                try:
-                    is_symlink = entry.is_symlink()
-                except OSError:
-                    # If is_symlink() raises an OSError, consider that the
-                    # entry is not a symbolic link, same behaviour than
-                    # os.path.islink().
-                    is_symlink = False
-                walk_into = not is_symlink
+                nondirs.append(entry.name)
 
-            if walk_into:
-                yield from walk(entry.path, topdown, onerror, followlinks)
+            if not topdown and is_dir:
+                # Bottom-up: recurse into sub-directory, but exclude symlinks to
+                # directories if followlinks is False
+                if followlinks:
+                    walk_into = True
+                else:
+                    try:
+                        is_symlink = entry.is_symlink()
+                    except OSError:
+                        # If is_symlink() raises an OSError, consider that the
+                        # entry is not a symbolic link, same behaviour than
+                        # os.path.islink().
+                        is_symlink = False
+                    walk_into = not is_symlink
+
+                if walk_into:
+                    yield from walk(entry.path, topdown, onerror, followlinks)
 
     # Yield before recursion if going top down
     if topdown:
@@ -437,15 +438,30 @@ class _DummyDirEntry:
     def __init__(self, dir, name):
         self.name = name
         self.path = path.join(dir, name)
+
     def is_dir(self):
         return path.isdir(self.path)
+
     def is_symlink(self):
         return path.islink(self.path)
 
-def _dummy_scandir(dir):
+class _dummy_scandir:
     # listdir-based implementation for bytes patches on Windows
-    for name in listdir(dir):
-        yield _DummyDirEntry(dir, name)
+    def __init__(self, dir):
+        self.dir = dir
+        self.it = iter(listdir(dir))
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        return _DummyDirEntry(self.dir, next(self.it))
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        self.it = iter(())
 
 __all__.append("walk")
 
diff --git a/Lib/test/test_os.py b/Lib/test/test_os.py
@@ -2808,6 +2808,8 @@ def test_os_all(self):
 
 
 class TestScandir(unittest.TestCase):
+    check_no_resource_warning = support.check_no_resource_warning
+
     def setUp(self):
         self.path = os.path.realpath(support.TESTFN)
         self.addCleanup(support.rmtree, self.path)
@@ -3030,6 +3032,56 @@ def test_bad_path_type(self):
         for obj in [1234, 1.234, {}, []]:
             self.assertRaises(TypeError, os.scandir, obj)
 
+    def test_close(self):
+        self.create_file("file.txt")
+        self.create_file("file2.txt")
+        iterator = os.scandir(self.path)
+        next(iterator)
+        iterator.close()
+        # multiple closes
+        iterator.close()
+        with self.check_no_resource_warning():
+            del iterator
+
+    def test_context_manager(self):
+        self.create_file("file.txt")
+        self.create_file("file2.txt")
+        with os.scandir(self.path) as iterator:
+            next(iterator)
+        with self.check_no_resource_warning():
+            del iterator
+
+    def test_context_manager_close(self):
+        self.create_file("file.txt")
+        self.create_file("file2.txt")
+        with os.scandir(self.path) as iterator:
+            next(iterator)
+            iterator.close()
+
+    def test_context_manager_exception(self):
+        self.create_file("file.txt")
+        self.create_file("file2.txt")
+        with self.assertRaises(ZeroDivisionError):
+            with os.scandir(self.path) as iterator:
+                next(iterator)
+                1/0
+        with self.check_no_resource_warning():
+            del iterator
+
+    def test_resource_warning(self):
+        self.create_file("file.txt")
+        self.create_file("file2.txt")
+        iterator = os.scandir(self.path)
+        next(iterator)
+        with self.assertWarns(ResourceWarning):
+            del iterator
+            support.gc_collect()
+        # exhausted iterator
+        iterator = os.scandir(self.path)
+        list(iterator)
+        with self.check_no_resource_warning():
+            del iterator
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/Misc/NEWS b/Misc/NEWS
@@ -179,6 +179,9 @@ Core and Builtins
 Library
 -------
 
+- Issue #25994: Added the close() method and the support of the context manager
+  protocol for the os.scandir() iterator.
+
 - Issue #23992: multiprocessing: make MapResult not fail-fast upon exception.
 
 - Issue #26243: Support keyword arguments to zlib.compress().  Patch by Aviv
diff --git a/Modules/posixmodule.c b/Modules/posixmodule.c
@@ -11937,8 +11937,14 @@ typedef struct {
 
 #ifdef MS_WINDOWS
 
+static int
+ScandirIterator_is_closed(ScandirIterator *iterator)
+{
+    return iterator->handle == INVALID_HANDLE_VALUE;
+}
+
 static void
-ScandirIterator_close(ScandirIterator *iterator)
+ScandirIterator_closedir(ScandirIterator *iterator)
 {
     if (iterator->handle == INVALID_HANDLE_VALUE)
         return;
@@ -11956,7 +11962,7 @@ ScandirIterator_iternext(ScandirIterator *iterator)
     BOOL success;
     PyObject *entry;
 
-    /* Happens if the iterator is iterated twice */
+    /* Happens if the iterator is iterated twice, or closed explicitly */
     if (iterator->handle == INVALID_HANDLE_VALUE)
         return NULL;
 
@@ -11987,14 +11993,20 @@ ScandirIterator_iternext(ScandirIterator *iterator)
     }
 
     /* Error or no more files */
-    ScandirIterator_close(iterator);
+    ScandirIterator_closedir(iterator);
     return NULL;
 }
 
 #else /* POSIX */
 
+static int
+ScandirIterator_is_closed(ScandirIterator *iterator)
+{
+    return !iterator->dirp;
+}
+
 static void
-ScandirIterator_close(ScandirIterator *iterator)
+ScandirIterator_closedir(ScandirIterator *iterator)
 {
     if (!iterator->dirp)
         return;
@@ -12014,7 +12026,7 @@ ScandirIterator_iternext(ScandirIterator *iterator)
     int is_dot;
     PyObject *entry;
 
-    /* Happens if the iterator is iterated twice */
+    /* Happens if the iterator is iterated twice, or closed explicitly */
     if (!iterator->dirp)
         return NULL;
 
@@ -12051,21 +12063,67 @@ ScandirIterator_iternext(ScandirIterator *iterator)
     }
 
     /* Error or no more files */
-    ScandirIterator_close(iterator);
+    ScandirIterator_closedir(iterator);
     return NULL;
 }
 
 #endif
 
+static PyObject *
+ScandirIterator_close(ScandirIterator *self, PyObject *args)
+{
+    ScandirIterator_closedir(self);
+    Py_RETURN_NONE;
+}
+
+static PyObject *
+ScandirIterator_enter(PyObject *self, PyObject *args)
+{
+    Py_INCREF(self);
+    return self;
+}
+
+static PyObject *
+ScandirIterator_exit(ScandirIterator *self, PyObject *args)
+{
+    ScandirIterator_closedir(self);
+    Py_RETURN_NONE;
+}
+
 static void
 ScandirIterator_dealloc(ScandirIterator *iterator)
 {
-    ScandirIterator_close(iterator);
+    if (!ScandirIterator_is_closed(iterator)) {
+        PyObject *exc, *val, *tb;
+        Py_ssize_t old_refcount = Py_REFCNT(iterator);
+        /* Py_INCREF/Py_DECREF cannot be used, because the refcount is
+         * likely zero, Py_DECREF would call again the destructor.
+         */
+        ++Py_REFCNT(iterator);
+        PyErr_Fetch(&exc, &val, &tb);
+        if (PyErr_WarnFormat(PyExc_ResourceWarning, 1,
+                             "unclosed scandir iterator %R", iterator)) {
+            /* Spurious errors can appear at shutdown */
+            if (PyErr_ExceptionMatches(PyExc_Warning))
+                PyErr_WriteUnraisable((PyObject *) iterator);
+        }
+        PyErr_Restore(exc, val, tb);
+        Py_REFCNT(iterator) = old_refcount;
+
+        ScandirIterator_closedir(iterator);
+    }
     Py_XDECREF(iterator->path.object);
     path_cleanup(&iterator->path);
     Py_TYPE(iterator)->tp_free((PyObject *)iterator);
 }
 
+static PyMethodDef ScandirIterator_methods[] = {
+    {"__enter__", (PyCFunction)ScandirIterator_enter, METH_NOARGS},
+    {"__exit__", (PyCFunction)ScandirIterator_exit, METH_VARARGS},
+    {"close", (PyCFunction)ScandirIterator_close, METH_NOARGS},
+    {NULL}
+};
+
 static PyTypeObject ScandirIteratorType = {
     PyVarObject_HEAD_INIT(NULL, 0)
     MODNAME ".ScandirIterator",             /* tp_name */
@@ -12095,6 +12153,7 @@ static PyTypeObject ScandirIteratorType = {
     0,                                      /* tp_weaklistoffset */
     PyObject_SelfIter,                      /* tp_iter */
     (iternextfunc)ScandirIterator_iternext, /* tp_iternext */
+    ScandirIterator_methods,                /* tp_methods */
 };
 
 static PyObject *
