Conditional use of os.walk

bids2table/_indexing.py

@@ -7,7 +7,6 @@
 import enum
 import fnmatch
 import importlib.metadata
-import os
 import re
 from concurrent.futures import Executor, ProcessPoolExecutor
 from functools import partial
@@ -22,7 +21,7 @@
     validate_bids_entities,
 )
 from ._logging import setup_logger
-from ._pathlib import PathT, as_path
+from ._pathlib import PathT, as_path, walk
 
 _BIDS_SUBJECT_DIR_PATTERN = re.compile(r"sub-[a-zA-Z0-9]+")
 
@@ -147,9 +146,9 @@ def find_bids_datasets(
     dir_count = 0
     ds_count = 0
 
-    for dirpath, dirnames, _ in os.walk(root, followlinks=follow_symlinks):
+    root = as_path(root)
+    for dirpath, dirnames, _ in walk(root, follow_symlinks=follow_symlinks):
         dir_count += 1
-        dirpath = as_path(dirpath)
 
         if _is_bids_dataset(dirpath):
             ds_count += 1

bids2table/_pathlib.py

@@ -1,4 +1,7 @@
+import os
+import sys
 from pathlib import Path
+from typing import Iterator
 
 try:
     from cloudpathlib import AnyPath, CloudPath, S3Client
@@ -18,6 +21,23 @@
 PathT = Path | CloudPath
 
 
+def walk(
+    root: str | PathT, follow_symlinks: bool = False
+) -> Iterator[tuple[str, list[str], list[str]]]:
+    if isinstance(root, str):
+        root = as_path(root)
+
+    # Py312+ or CloudPath
+    if sys.version_info >= (3, 12) or (
+        _CLOUDPATHLIB_AVAILABLE and isinstance(root, CloudPath)
+    ):
+        yield from root.walk(follow_symlinks=follow_symlinks)
+    # Fall back to os.walk for local paths
+    else:
+        for dirpath, dirnames, filenames in os.walk(root, followlinks=follow_symlinks):
+            yield Path(dirpath), dirnames, filenames
+
+
 def as_path(path: str | PathT) -> PathT:
     """Cast input to a `Path` type."""
     if isinstance(path, str):