From ffe96ae10be8a3117fa18c35034fcfc45c3cf7b7 Mon Sep 17 00:00:00 2001 From: Serhiy Storchaka Date: Thu, 11 Feb 2016 13:21:30 +0200 Subject: [PATCH] Issue #25994: Added the close() method and the support of the context manager protocol for the os.scandir() iterator. --- Doc/library/os.rst | 27 ++++++++++-- Doc/whatsnew/3.6.rst | 11 +++++ Lib/os.py | 96 +++++++++++++++++++++++++------------------ Lib/test/test_os.py | 52 +++++++++++++++++++++++ Misc/NEWS | 3 ++ Modules/posixmodule.c | 73 ++++++++++++++++++++++++++++---- 6 files changed, 212 insertions(+), 50 deletions(-) diff --git a/Doc/library/os.rst b/Doc/library/os.rst index f16a3fbeca7..f064717bb0e 100644 --- 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 index bc5e5502752..b7d14f29526 100644 --- 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 index 674a7d7efda..c3f674ec3d1 100644 --- 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) - - 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 + entry = next(scandir_it) + except StopIteration: + break + except OSError as error: + if onerror is not None: + onerror(error) + return - if walk_into: - yield from walk(entry.path, topdown, onerror, followlinks) + 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) + + 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 index 66a426a4f13..07682f2bf9a 100644 --- a/Lib/test/test_os.py +++ b/Lib/test/test_os.py @@ -2808,6 +2808,8 @@ class ExportsTests(unittest.TestCase): 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 @@ class TestScandir(unittest.TestCase): 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 index 405e0b986d3..92a210fe121 100644 --- 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 index 2688cbcc23a..6e0081d0dff 100644 --- 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 *