Repository: google/etils

Issue:
Support walk in epath?
hi there!

is there a plan to support `walk` in epath for different backends? i can contribute if this feels like a good idea.


Discussion:
Hi, yes, this sounds like a good addition for parity with `pathlib`. Feel free to send a PR, happy to review :)

Resolving patch:
diff --git a/etils/epath/abstract_path.py b/etils/epath/abstract_path.py
--- a/etils/epath/abstract_path.py
+++ b/etils/epath/abstract_path.py
@@ -19,7 +19,7 @@
 import os
 import pathlib
 import typing
-from typing import Any, AnyStr, Iterator, Optional, Type, TypeVar
+from typing import Any, AnyStr, Iterator, Optional, Type, TypeVar, Callable
 
 from etils.epath import register
 from etils.epath import stat_utils
@@ -112,6 +112,15 @@ def rglob(self: _T, pattern: str) -> Iterator[_T]:
     """Yields all matching files recursively (of any kind)."""
     return self.glob(f'**/{pattern}')
 
+  @abstractmethod
+  def walk(
+      self: _T,
+      top_down: bool = True,
+      on_error: Callable[[OSError], object] | None = None,
+      follow_symlinks: bool = False,
+  ) -> Iterator[tuple[_T, list[str], list[str]]]:
+    raise NotImplementedError
+
   def expanduser(self: _T) -> _T:
     """Returns a new path with expanded `~` and `~user` constructs."""
     if '~' not in self.parts:  # pytype: disable=attribute-error
@@ -218,4 +227,5 @@ def replace(self: _T, target: PathLike) -> _T:
   @abstractmethod
   def copy(self: _T, dst: PathLike, overwrite: bool = False) -> _T:
     """Copy the current file to the given destination."""
+
   # pytype: enable=bad-return-type
diff --git a/etils/epath/backend.py b/etils/epath/backend.py
--- a/etils/epath/backend.py
+++ b/etils/epath/backend.py
@@ -24,8 +24,7 @@
 import shutil
 import stat as stat_lib
 import typing
-from typing import Iterator, NoReturn, Optional, Union
-
+from typing import Iterator, NoReturn, Optional, Union, Callable
 from etils.epath import stat_utils
 from etils.epath.typing import PathLike  # pylint: disable=g-importing-member
 
@@ -105,6 +104,16 @@ def copy(self, path: PathLike, dst: PathLike, overwrite: bool) -> None:
   def stat(self, path: PathLike) -> stat_utils.StatResult:
     raise NotImplementedError
 
+  @abc.abstractmethod
+  def walk(
+      self,
+      top: PathLike,
+      top_down: bool = True,
+      on_error: Callable[[OSError], object] | None  = None,
+      follow_symlinks: bool = False,
+  ) -> Iterator[tuple[PathLike, list[str], list[str]]]:
+    raise NotImplementedError
+
 
 class _OsPathBackend(Backend):
   """`os.path` backend."""
@@ -218,6 +227,23 @@ def stat(self, path: PathLike) -> stat_utils.StatResult:
         mode=st.st_mode,
     )
 
+  def walk(
+      self,
+      top: PathLike,
+      top_down: bool = True,
+      on_error: Callable[[OSError], object] | None = None,
+      follow_symlinks: bool = False,
+  ) -> Iterator[tuple[PathLike, list[str], list[str]]]:
+    for path, dirnames, filenames in os.walk(
+        top, topdown=top_down, onerror=on_error, followlinks=follow_symlinks
+    ):
+      #  Unlike os.walk(), Path.walk() lists symlinks to directories in filenames if follow_symlinks is false
+      if follow_symlinks is False:
+        symlinks = {dirname for dirname in dirnames if os.path.islink(os.path.join(path, dirname))}
+        dirnames[:] = list(dirname for dirname in dirnames if dirname not in symlinks) # we need to keep the reference to the original dirnames object
+        filenames.extend(symlinks)
+      yield path, dirnames, filenames
+
 
 class _TfBackend(Backend):
   """TensorFlow backend."""
@@ -381,6 +407,17 @@ def stat(self, path: PathLike) -> stat_utils.StatResult:
         mode=None,
     )
 
+  def walk(
+      self,
+      top: PathLike,
+      top_down: bool = True,
+      on_error: Callable[[OSError], object] | None = None,
+      follow_symlinks: bool = False,
+  ) -> Iterator[tuple[PathLike, list[str], list[str]]]:
+    yield from self.gfile.walk(
+        top, topdown=top_down, onerror=on_error
+    )
+
 
 class _FileSystemSpecBackend(Backend):
   """FileSystemSpec backend entirely relying on fsspec."""
@@ -538,7 +575,21 @@ def stat(self, path: PathLike) -> stat_utils.StatResult:
         mtime=mtime,
         owner=info.get('owner'),
         group=info.get('group'),
-        mode=info.get('mode')
+        mode=info.get('mode'),
+    )
+
+  def walk(
+      self,
+      top: PathLike,
+      top_down: bool = True,
+      on_error: Callable[[OSError], object] | None = None,
+      follow_symlinks: bool = False,
+  ) -> Iterator[tuple[PathLike, list[str], list[str]]]:
+
+    if on_error is None:
+      on_error = 'omit'  # default behavior for pathlib.Path.walk
+    yield from self.fs(top).walk(
+        path=top, topdown=top_down, on_error=on_error, max_depth=None
     )
 
 
diff --git a/etils/epath/gpath.py b/etils/epath/gpath.py
--- a/etils/epath/gpath.py
+++ b/etils/epath/gpath.py
@@ -24,7 +24,7 @@
 import posixpath
 import types
 import typing
-from typing import Any, ClassVar, Iterator, Optional, Type, TypeVar, Union
+from typing import Any, ClassVar, Iterator, Optional, Type, TypeVar, Union, Callable, Iterable
 
 from etils import epy
 from etils.epath import abstract_path
@@ -270,6 +270,20 @@ def stat(self) -> stat_utils.StatResult:
     """Returns metadata for the file/directory."""
     return self._backend.stat(self._path_str)
 
+  def walk(
+      self: _P,
+      top_down: bool = True,
+      on_error: Callable[[OSError], object] | None  = None,
+      follow_symlinks: bool = False,
+  ) -> Iterator[tuple[_P, list[str], list[str]]]:
+    for root, dirs, files in self._backend.walk(
+        top=self._path_str,
+        top_down=top_down,
+        on_error=on_error,
+        follow_symlinks=follow_symlinks,
+    ):
+      yield self._new(root), dirs, files
+
 
 def _get_backend(p0: _GPath, p1: _GPath) -> backend_lib.Backend:
   """When composing with another backend, GCS win.


Test patch:
diff --git a/etils/epath/backend_test.py b/etils/epath/backend_test.py
--- a/etils/epath/backend_test.py
+++ b/etils/epath/backend_test.py
@@ -150,6 +150,47 @@ def _test_glob(
   ]
 
 
+def _test_walk(backend: epath.backend.Backend, tmp_path):
+  nested = (tmp_path / 'abc/nested/')
+  nested.mkdir(parents=True)
+  other_nested = (tmp_path / 'abc/other_nested/')
+  other_nested.mkdir(parents=True)
+  (nested / '001').touch()
+  (nested / '002').touch()
+  (nested / '003').touch()
+  linked_dir = (tmp_path / 'abc/link_dir')
+  linked_dir.symlink_to(nested)
+
+  bottom_up_walk = list(backend.walk(tmp_path, top_down=False, follow_symlinks=False))
+
+
+  assert bottom_up_walk == [
+    (str(other_nested), [], []),
+    (str(nested), [], ['003', '002', '001']),
+    (str(tmp_path / 'abc'), [ 'other_nested', 'nested'], ['link_dir']),
+    (str(tmp_path), ['abc'], [])
+  ]
+
+  top_down_walk = list(backend.walk(tmp_path, top_down=True, follow_symlinks=False))
+
+  assert top_down_walk == [
+    (str(tmp_path), ['abc'], []),
+    (str(tmp_path / 'abc'), [ 'other_nested', 'nested'], ['link_dir']),
+    (str(other_nested), [], []),
+    (str(nested), [], ['003', '002', '001']),
+  ]
+
+  if isinstance(backend, epath.backend._OsPathBackend):
+    follow_symlinks_walk = list(backend.walk(tmp_path, follow_symlinks=True))
+    assert follow_symlinks_walk == [
+      (str(tmp_path), ['abc'], []),
+      (str(tmp_path / 'abc'), ['link_dir', 'other_nested', 'nested'], []),
+      (str(linked_dir), [], ['003', '002', '001']),
+      (str(other_nested), [], []),
+      (str(nested), [], ['003', '002', '001']),
+    ]
+
+
 def _test_listdir(
     backend: epath.backend.Backend,
     tmp_path: pathlib.Path,
@@ -513,6 +554,7 @@ def _test_stat(
         _test_copy,
         _test_copy_with_overwrite,
         _test_stat,
+        _test_walk,
     ],
 )
 def test_backend(
diff --git a/etils/epath/gpath_test.py b/etils/epath/gpath_test.py
--- a/etils/epath/gpath_test.py
+++ b/etils/epath/gpath_test.py
@@ -69,6 +69,7 @@ def new_fn(_, p, **kwargs):
       mkdir=_call(backend.mkdir),
       remove=_call(backend.remove),
       rmtree=_call(backend.rmtree),
+      walk=_call(backend.walk),
       # Mock _is_tf_installed in order not to default to tf_backend:
   ), mock.patch.object(epath.gpath, '_is_tf_installed', return_value=False):
     yield tmp_path
@@ -398,6 +399,22 @@ def test_rglob():
   with pytest.raises(NotImplementedError, match='Recursive'):
     list(epath.Path('/tmp/nonexisting/test').glob('folder/**/img.jpg'))
 
+def test_walk(gcs_mocked_path: pathlib.Path):
+  root_path = epath.Path(gcs_mocked_path)
+  result = list((root_path / '/nonexisting/test').walk())
+  assert len(result) == 0
+
+  subdir = (root_path / 'subdir')
+  subdir.mkdir()
+  (root_path / '001').touch()
+  (root_path / '002').touch()
+  (root_path / '003').touch()
+
+  result = list(root_path.walk(top_down=True, follow_symlinks=False))
+  assert result == [
+    (root_path, ['subdir'], ['003', '002', '001']),
+    (subdir, [], [])
+  ]
 
 def test_default():
   path = epath.Path()
diff --git a/etils/epath/testing.py b/etils/epath/testing.py
--- a/etils/epath/testing.py
+++ b/etils/epath/testing.py
@@ -123,6 +123,7 @@ def mock_epath(
     replace: Optional[_MockFn] = None,
     rmtree: Optional[_MockFn] = None,
     stat: Optional[_MockFn] = None,
+    walk: Optional[_MockFn] = None,
 ) -> Iterator[None]:
   """Mock epath.
 
@@ -163,7 +164,7 @@ def mock_epath(
       replace=replace,
       rmtree=rmtree,
       stat=stat,
-      # 'walk',
+      walk = walk,
   )
   mock_backend = _MockBackend(mock_fns=mock_fns)