Prefer scandir-based gather for source-tree scanning

mkpfs/gather.py

@@ -0,0 +1,65 @@
+"""Gather helpers for MkPFS.
+
+Provide a scandir-based file gather implementation as a faster, low-level
+alternative to Path.rglob("*"). Keep this module minimal and dependency-free
+other than mkpfs.utils.is_ignored_name.
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+from .utils import is_ignored_name
+
+
+def gather_files_scandir(root: Path) -> list[Path]:
+    """Walk ``root`` using os.scandir and return a list of file Paths.
+
+    Rules:
+    - Skip names where ``is_ignored_name(name)`` is True (applies to files and
+      directories).
+    - Do NOT follow directory symlinks (to avoid loops). File symlinks are
+      included (entry.is_file() follows symlinks).
+    - Return Path objects (not resolved); the caller is responsible for any
+      normalization or resolution.
+
+    The caller is responsible for any additional validation (non-ASCII names,
+    deterministic global sorting, etc.). This helper focuses on fast traversal.
+    """
+    root = Path(root)
+    abs_files: list[Path] = []
+    stack: list[Path] = [root]
+
+    while stack:
+        dir_path = stack.pop()
+        try:
+            with os.scandir(dir_path) as it:
+                entries = list(it)
+        except PermissionError:
+            # Mirror conservative behavior: skip directories we cannot access.
+            continue
+
+        # Deterministic per-directory ordering to reduce nondeterminism while
+        # keeping traversal fast. The final global sort is left to the caller.
+        entries.sort(key=lambda e: e.name.lower())
+
+        for entry in entries:
+            name = entry.name
+            if is_ignored_name(name):
+                continue
+
+            try:
+                # Directories: descend but do NOT follow symlinks
+                if entry.is_dir(follow_symlinks=False):
+                    stack.append(Path(dir_path) / name)
+                    continue
+
+                # Files: include regular files and file symlinks
+                if entry.is_file():
+                    abs_files.append(Path(dir_path) / name)
+            except OSError:
+                # If the entry disappears or cannot be inspected, skip it.
+                continue
+
+    return abs_files

mkpfs/pfs.py

@@ -35,6 +35,7 @@
 from . import consts
 from .exfat import EXFAT_SIGNATURE, ExfatEntry, ExfatError, ExfatReader
 from .exfat_writer import iter_exfat_image
+from .gather import gather_files_scandir
 from .logging import info, warning
 from .pbar import Progress
 from .utils import (
@@ -2695,11 +2696,10 @@ def scan_source_tree(root: Path, progress: Progress) -> tuple[dict[str, DirNode]
     progress.status("\nDiscovering files...")
     # Exclude OS-generated metadata (.DS_Store, ._*, Thumbs.db, __MACOSX, ...) so it
     # never ends up in the image, including files nested under an ignored directory.
-    abs_files: list[Path] = [
-        p
-        for p in root.rglob("*")
-        if p.is_file() and not any(is_ignored_name(part) for part in p.relative_to(root).parts)
-    ]
+    # Prefer scandir-based gather for performance on large trees. Default is scandir.
+    abs_files = gather_files_scandir(root)
+    # Keep safety check: re-apply ignore filter in case helper included unexpected entries
+    abs_files = [p for p in abs_files if not any(is_ignored_name(part) for part in p.relative_to(root).parts)]
     abs_files.sort(key=lambda p: p.relative_to(root).as_posix().lower())
 
     # Validate filenames before compression work begins; non-ASCII names are unsupported.

tests/mkpfs/test_gather.py

@@ -0,0 +1,59 @@
+from pathlib import Path
+
+import pytest
+
+from mkpfs.gather import gather_files_scandir
+from mkpfs.pbar import Progress
+from mkpfs.pfs import BuildError, scan_source_tree
+from mkpfs.utils import is_ignored_name
+
+
+def _sorted_rel_paths(root: Path, paths: list[Path]) -> list[str]:
+    return sorted([p.relative_to(root).as_posix() for p in paths])
+
+
+def test_scandir_equals_rglob_small_tree(tmp_path: Path) -> None:
+    # Create tree
+    (tmp_path / "a").mkdir()
+    (tmp_path / "a" / "file1.txt").write_text("x")
+    (tmp_path / "b").mkdir()
+    (tmp_path / "b" / "file2.txt").write_text("y")
+    (tmp_path / "b" / "sub").mkdir()
+    (tmp_path / "b" / "sub" / "file3.txt").write_text("z")
+
+    # rglob baseline
+    rglob_files = [
+        p
+        for p in tmp_path.rglob("*")
+        if p.is_file() and not any(is_ignored_name(part) for part in p.relative_to(tmp_path).parts)
+    ]
+    rglob_sorted = _sorted_rel_paths(tmp_path, rglob_files)
+
+    # scandir gather
+    scandir_files = gather_files_scandir(tmp_path)
+    scandir_sorted = _sorted_rel_paths(tmp_path, scandir_files)
+
+    assert rglob_sorted == scandir_sorted
+
+
+def test_ignored_names_are_excluded(tmp_path: Path) -> None:
+    (tmp_path / "a").mkdir()
+    (tmp_path / "a" / ".DS_Store").write_text("x")
+    (tmp_path / "b").mkdir()
+    (tmp_path / "b" / "file.txt").write_text("y")
+
+    scandir_files = gather_files_scandir(tmp_path)
+    rels = [p.relative_to(tmp_path).as_posix() for p in scandir_files]
+    assert ".DS_Store" not in rels
+    assert "b/file.txt" in rels
+
+
+def test_scan_source_tree_non_ascii_raises(tmp_path: Path) -> None:
+    # Create a file with non-ascii name
+    (tmp_path / "dir").mkdir()
+    bad_name = "não.txt"
+    (tmp_path / "dir" / bad_name).write_text("x")
+
+    # scan_source_tree should detect non-ascii names and raise BuildError
+    with pytest.raises(BuildError):
+        scan_source_tree(tmp_path, progress=Progress(enabled=False))