Return context-manager iterator from scandir() (#363)

Abandoning scandir() iteration before exhaustion left the SMB
directory handle open until GC finalized the generator, and
callers had no idiomatic way to release it sooner. Stdlib
os.scandir() solves this by returning a context-manager
iterator.

Author: henrik-nil-acc

src/smbclient/__init__.py

@@ -8,6 +8,7 @@
     XATTR_REPLACE,
     SMBDirEntry,
     SMBDirEntryInformation,
+    SMBScandirIterator,
     SMBStatResult,
     SMBStatVolumeResult,
     copyfile,

src/smbclient/_os.py

@@ -666,23 +666,62 @@ def rmdir(path, **kwargs):
     _delete(SMBDirectoryIO, path, **kwargs)
 
 
-def scandir(path, search_pattern="*", **kwargs):
+class SMBScandirIterator:
+    """Iterator over SMB directory entries with ``with``-driven close.
+
+    Iterable directly and usable as a context manager whose exit releases
+    the SMB directory handle.
+    """
+
+    __slots__ = ("_gen",)
+
+    def __init__(self, gen: t.Generator[SMBDirEntry, None, None]) -> None:
+        self._gen = gen
+
+    def __iter__(self) -> SMBScandirIterator:
+        return self
+
+    def __next__(self) -> SMBDirEntry:
+        return next(self._gen)
+
+    def __enter__(self) -> SMBScandirIterator:
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        self.close()
+
+    def close(self) -> None:
+        self._gen.close()
+
+
+def scandir(path: str, search_pattern: str = "*", **kwargs: t.Any) -> SMBScandirIterator:
     """
     Return an iterator of DirEntry objects corresponding to the entries in the directory given by path. The entries are
     yielded in arbitrary order, and the special entries '.' and '..' are not included.
 
+    Mirrors stdlib ``os.scandir()``: the returned iterator also supports the context-manager protocol so callers can
+    release the SMB directory handle deterministically:
+
+        with smbclient.scandir(path) as it:
+            for entry in it:
+                ...
+
     Using scandir() instead of listdir() can significantly increase the performance of code that also needs file type
     or file attribute information, because DirEntry objects expose this information if the SMB server provides it when
     scanning a directory. All DirEntry methods may perform a SMB request, but is_dir(), is_file(), is_symlink() usually
     only require a one system call unless the file or directory is a reparse point which requires 2 calls. See the
     Python documentation for how DirEntry is set up and the methods and attributes that are available.
 
     :param path: The path to a directory to scan.
-    :param search_pattern: THe search string to match against the names of directories or files. This pattern can use
+    :param search_pattern: The search string to match against the names of directories or files. This pattern can use
     '*' as a wildcard for multiple chars and '?' as a wildcard for a single char. Does not support regex patterns.
     :param kwargs: Common SMB Session arguments for smbclient.
-    :return: An iterator of DirEntry objects in the directory.
+    :return: A context-manager iterator of DirEntry objects in the directory.
     """
+    return SMBScandirIterator(_scandir(path, search_pattern, **kwargs))
+
+
+def _scandir(path: str, search_pattern: str = "*", **kwargs: t.Any) -> t.Generator[SMBDirEntry, None, None]:
     connection_cache = kwargs.get("connection_cache", None)
     with SMBDirectoryIO(path, share_access="rwd", **kwargs) as fd:
         for raw_dir_info in fd.query_directory(search_pattern, FileInformationClass.FILE_ID_FULL_DIRECTORY_INFORMATION):
@@ -1023,26 +1062,27 @@ def walk(top, topdown=True, onerror=None, follow_symlinks=False, **kwargs):
     dirs = []
     files = []
     bottom_up_dirs = []
-    while True:
-        try:
+    with scandir_gen:
+        while True:
             try:
-                entry = next(scandir_gen)
-            except StopIteration:
-                break
-        except OSError as err:
-            if onerror is not None:
-                onerror(err)
-            return
-
-        if not entry.is_dir():
-            files.append(entry.name)
-            continue
+                try:
+                    entry = next(scandir_gen)
+                except StopIteration:
+                    break
+            except OSError as err:
+                if onerror is not None:
+                    onerror(err)
+                return
+
+            if not entry.is_dir():
+                files.append(entry.name)
+                continue
 
-        dirs.append(entry.name)
-        if not topdown and (follow_symlinks or not entry.is_symlink()):
-            # Add the directory to the bottom up list which is recursively walked below, we exclude symlink dirs if
-            # follow_symlinks is False.
-            bottom_up_dirs.append(entry.path)
+            dirs.append(entry.name)
+            if not topdown and (follow_symlinks or not entry.is_symlink()):
+                # Add the directory to the bottom up list which is recursively walked below, we exclude symlink dirs
+                # if follow_symlinks is False.
+                bottom_up_dirs.append(entry.path)
 
     walk_kwargs = {"topdown": topdown, "onerror": onerror, "follow_symlinks": follow_symlinks}
     walk_kwargs.update(kwargs)

src/smbclient/shutil.py

@@ -310,7 +310,8 @@ def copytree(
     :return: The dst path.
     """
     if is_remote_path(src):
-        dir_entries = list(scandir(src, **kwargs))
+        with scandir(src, **kwargs) as scandir_gen:
+            dir_entries = list(scandir_gen)
     else:
         dir_entries = list(os.scandir(src))
 
@@ -419,34 +420,34 @@ def onerror(*args):
         onerror(islink, path, sys.exc_info())
         return
 
-    scandir_gen = scandir(path, **kwargs)
-    while True:
-        try:
-            dir_entry = next(scandir_gen)
-        except StopIteration:
-            break
-        except OSError:
-            onerror(scandir, path, sys.exc_info())
-            continue
-
-        # In case the entry is a directory symbolic link we need to remove the dir itself and not recurse down into
-        # it with rmtree. Doing that would result in a symbolic link target having it's contents removed even if it's
-        # outside the rmtree scope.
-        if (
-            dir_entry.is_symlink()
-            and dir_entry.stat(follow_symlinks=False).st_file_attributes & FileAttributes.FILE_ATTRIBUTE_DIRECTORY
-        ):
+    with scandir(path, **kwargs) as scandir_gen:
+        while True:
             try:
-                rmdir(dir_entry.path, **kwargs)
+                dir_entry = next(scandir_gen)
+            except StopIteration:
+                break
             except OSError:
-                onerror(rmdir, dir_entry.path, sys.exc_info())
-        elif dir_entry.is_dir():
-            rmtree(dir_entry.path, ignore_errors, onerror, **kwargs)
-        else:
-            try:
-                remove(dir_entry.path, **kwargs)
-            except OSError:
-                onerror(remove, dir_entry.path, sys.exc_info())
+                onerror(scandir, path, sys.exc_info())
+                continue
+
+            # In case the entry is a directory symbolic link we need to remove the dir itself and not recurse down into
+            # it with rmtree. Doing that would result in a symbolic link target having it's contents removed even if
+            # it's outside the rmtree scope.
+            if (
+                dir_entry.is_symlink()
+                and dir_entry.stat(follow_symlinks=False).st_file_attributes & FileAttributes.FILE_ATTRIBUTE_DIRECTORY
+            ):
+                try:
+                    rmdir(dir_entry.path, **kwargs)
+                except OSError:
+                    onerror(rmdir, dir_entry.path, sys.exc_info())
+            elif dir_entry.is_dir():
+                rmtree(dir_entry.path, ignore_errors, onerror, **kwargs)
+            else:
+                try:
+                    remove(dir_entry.path, **kwargs)
+                except OSError:
+                    onerror(remove, dir_entry.path, sys.exc_info())
 
     try:
         rmdir(path, **kwargs)

tests/conftest.py

@@ -154,6 +154,32 @@
 )
 
 
+class StubScandirGen:
+    """Server-free stand-in for the _scandir() generator: injects a
+    mid-iteration failure and records close() calls, so the caller
+    close-on-exception tests need no SMB server."""
+
+    def __init__(self, closes, next_exc=None):
+        self._closes = closes
+        self._next_exc = next_exc
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        if self._next_exc is not None:
+            raise self._next_exc
+        raise StopIteration
+
+    def close(self):
+        self._closes.append(True)
+
+
+@pytest.fixture
+def stub_scandir_gen():
+    return StubScandirGen
+
+
 @pytest.fixture(scope="module")
 def smb_real():
     # for these tests to work the server at SMB_SERVER must support dialect

tests/test_smbclient_os.py

@@ -1308,6 +1308,36 @@ def test_scandir_with_non_matching_pattern(smb_share):
     assert list(smbclient.scandir(smb_share, search_pattern="nomatch_*")) == []
 
 
+def test_scandir_as_context_manager(smb_share):
+    for filename in ["file1.txt", "file2.txt"]:
+        with smbclient.open_file(rf"{smb_share}\{filename}", mode="w") as fd:
+            fd.write("content")
+
+    # Full iteration inside the context manager yields every entry.
+    with smbclient.scandir(smb_share) as scandir_gen:
+        assert isinstance(scandir_gen, smbclient.SMBScandirIterator)
+        assert sorted(entry.name for entry in scandir_gen) == ["file1.txt", "file2.txt"]
+
+    # Abandoning iteration early still releases the handle on block exit: the
+    # iterator is finalised, so resuming it stops instead of yielding the rest.
+    it = smbclient.scandir(smb_share)
+    with it:
+        assert next(it).name in ("file1.txt", "file2.txt")
+    with pytest.raises(StopIteration):
+        next(it)
+
+
+def test_scandir_iterator_contract():
+    it = smbclient.SMBScandirIterator(x for x in ["a", "b"])
+    with it as entered:
+        assert entered is it
+        assert next(it) == "a"
+
+    # __exit__ closed the underlying generator, so iteration is exhausted.
+    with pytest.raises(StopIteration):
+        next(it)
+
+
 @pytest.mark.skipif(
     os.name != "nt" and not os.environ.get("SMB_FORCE", False), reason="cannot create symlinks on Samba"
 )
@@ -1997,6 +2027,19 @@ def test_walk_with_symlink_dont_follow(smb_share):
     assert scanned_roots[src_dirname]["files"] == ["file.txt"]
 
 
+def test_walk_closes_scandir_iterator_on_unhandled_exception(monkeypatch, stub_scandir_gen):
+    closes = []
+    monkeypatch.setattr(
+        "smbclient._os._scandir",
+        lambda *a, **kw: stub_scandir_gen(closes, RuntimeError("simulated mid-iter failure")),
+    )
+
+    with pytest.raises(RuntimeError, match="simulated mid-iter failure"):
+        list(smbclient.walk(r"\\server\share\dir"))
+
+    assert closes == [True]
+
+
 def test_xattr_file(smb_share):
     filename = "%s\\file.txt" % smb_share
 

tests/test_smbclient_shutil.py

@@ -1607,3 +1607,30 @@ def _failing_islink(*args, **kwargs):
     assert callback_args[0][0].__name__ == "islink"
     assert callback_args[0][1] == fake_path
     assert isinstance(callback_args[0][2][1], SMBOSError)
+
+
+def test_rmtree_closes_scandir_iterator_on_unhandled_exception(monkeypatch, stub_scandir_gen):
+    closes = []
+    monkeypatch.setattr(
+        "smbclient._os._scandir",
+        lambda *a, **kw: stub_scandir_gen(closes, RuntimeError("simulated mid-iter failure")),
+    )
+    monkeypatch.setattr("smbclient.shutil.islink", lambda *a, **kw: False)
+
+    with pytest.raises(RuntimeError, match="simulated mid-iter failure"):
+        rmtree(r"\\server\share\dst")
+
+    assert closes == [True]
+
+
+def test_copytree_closes_scandir_iterator_on_unhandled_exception(monkeypatch, stub_scandir_gen):
+    closes = []
+    monkeypatch.setattr(
+        "smbclient._os._scandir",
+        lambda *a, **kw: stub_scandir_gen(closes, RuntimeError("simulated mid-iter failure")),
+    )
+
+    with pytest.raises(RuntimeError, match="simulated mid-iter failure"):
+        copytree(r"\\server\share\src", r"\\server\share\dst")
+
+    assert closes == [True]