magic patch os.listdir (#695)

* magic patch os.listdir

* Update node.py

* remove io patch

tests/integration/test_fs_patch.py

@@ -1,3 +1,5 @@
+import os
+
 import git
 
 import zntrack.examples
@@ -19,9 +21,33 @@ def test_patch_open(proj_path):
 
     node = zntrack.examples.WriteDVCOuts.from_rev(rev=commit.hexsha)
 
-    with node.state.patch_open():
+    with node.state.magic_patch():
         with open(node.outs, "r") as f:
             assert f.read() == "Hello World"
 
     with open(node.outs, "r") as f:
         assert f.read() == "Lorem Ipsum"
+
+    listdir = os.listdir(node.nwd)
+
+    with node.state.magic_patch():
+        os.listdir(node.nwd) == listdir
+
+
+def test_patch_list(proj_path):
+    node = zntrack.from_rev(
+        "HelloWorld",
+        remote="https://github.com/PythonFZ/ZnTrackExamples.git",
+        rev="fbb6ada",
+    )
+
+    def func(self, path):
+        return os.listdir(path)
+
+    type(node).list = func
+
+    with node.state.magic_patch():
+        assert "nodes/HelloWorld/random_number.json" in node.list(node.nwd)
+        assert "nodes/HelloWorld/node-meta.json" in node.list(node.nwd)
+
+    assert node.list(node.nwd) == []

zntrack/core/node.py

@@ -6,6 +6,7 @@
 import functools
 import json
 import logging
+import os
 import pathlib
 import time
 import typing
@@ -77,13 +78,14 @@ def fs(self) -> dvc.api.DVCFileSystem:
         raise dvc.utils.strictyaml.YAMLValidationError
 
     @contextlib.contextmanager
-    def patch_open(self) -> typing.ContextManager:
+    def magic_patch(self) -> typing.ContextManager:
         """Patch the open function to use the Node's file system.
 
         Opening a relative path will use the Node's file system.
         Opening an absolute path will use the local file system.
         """
         original_open = open
+        original_listdir = os.listdir
 
         def _open(file, *args, **kwargs):
             if file == "params.yaml":
@@ -94,10 +96,17 @@ def _open(file, *args, **kwargs):
 
             return original_open(file, *args, **kwargs)
 
+        def _listdir(path, *args, **kwargs):
+            if not pathlib.Path(path).is_absolute():
+                return self.fs.listdir(path, detail=False)
+
+            return original_listdir(path, *args, **kwargs)
+
         with unittest.mock.patch("builtins.open", _open):
             with unittest.mock.patch("__main__.open", _open):
-                # Jupyter Notebooks replace open with io.open
-                yield
+                with unittest.mock.patch("os.listdir", _listdir):
+                    # Jupyter Notebooks replace open with io.open
+                    yield
 
 
 class _NameDescriptor(zninit.Descriptor):