From e8cb07b8f237902370250f8621b4dc475a2e9668 Mon Sep 17 00:00:00 2001 From: Henrik Nilsson Date: Tue, 19 May 2026 04:47:06 +0000 Subject: [PATCH] Return context-manager iterator from scandir() 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. --- src/smbclient/__init__.py | 1 + src/smbclient/_os.py | 82 +++++++++++++++++++++++++--------- src/smbclient/shutil.py | 55 ++++++++++++----------- tests/conftest.py | 26 +++++++++++ tests/test_smbclient_os.py | 43 ++++++++++++++++++ tests/test_smbclient_shutil.py | 27 +++++++++++ 6 files changed, 186 insertions(+), 48 deletions(-) diff --git a/src/smbclient/__init__.py b/src/smbclient/__init__.py index a6d61d2..4e11a24 100644 --- a/src/smbclient/__init__.py +++ b/src/smbclient/__init__.py @@ -8,6 +8,7 @@ XATTR_REPLACE, SMBDirEntry, SMBDirEntryInformation, + SMBScandirIterator, SMBStatResult, SMBStatVolumeResult, copyfile, diff --git a/src/smbclient/_os.py b/src/smbclient/_os.py index f85ed41..a318521 100644 --- a/src/smbclient/_os.py +++ b/src/smbclient/_os.py @@ -666,11 +666,46 @@ 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 @@ -678,11 +713,15 @@ def scandir(path, search_pattern="*", **kwargs): 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) diff --git a/src/smbclient/shutil.py b/src/smbclient/shutil.py index 43e3d48..6a50850 100644 --- a/src/smbclient/shutil.py +++ b/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) diff --git a/tests/conftest.py b/tests/conftest.py index bf63501..4c0814c 100644 --- a/tests/conftest.py +++ b/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 diff --git a/tests/test_smbclient_os.py b/tests/test_smbclient_os.py index e1387e2..44fabfc 100644 --- a/tests/test_smbclient_os.py +++ b/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 diff --git a/tests/test_smbclient_shutil.py b/tests/test_smbclient_shutil.py index ca3d604..d8c2be3 100644 --- a/tests/test_smbclient_shutil.py +++ b/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]