patch & test pathlib

tests/integration/test_fs_patch.py → tests/integration/test_magic_patch.py

@@ -1,4 +1,5 @@
 import os
+import pathlib
 
 import git
 
@@ -24,9 +25,15 @@ def test_patch_open(proj_path):
     with node.state.magic_patch():
         with open(node.outs, "r") as f:
             assert f.read() == "Hello World"
+        assert pathlib.Path(node.outs).read_text() == "Hello World"
+        with pathlib.Path(node.outs).open("r") as f:
+            assert f.read() == "Hello World"
 
     with open(node.outs, "r") as f:
         assert f.read() == "Lorem Ipsum"
+    assert pathlib.Path(node.outs).read_text() == "Lorem Ipsum"
+    with pathlib.Path(node.outs).open("r") as f:
+        assert f.read() == "Lorem Ipsum"
 
     listdir = os.listdir(node.nwd)
 

zntrack/core/node.py

@@ -72,7 +72,9 @@ def magic_patch(self) -> typing.ContextManager:
         Opening a relative path will use the Node's file system.
         Opening an absolute path will use the local file system.
         """
-        original_open = open
+        import builtins
+
+        original_open = builtins.open
         original_listdir = os.listdir
 
         def _open(file, *args, **kwargs):
@@ -84,6 +86,17 @@ def _open(file, *args, **kwargs):
 
             return original_open(file, *args, **kwargs)
 
+        def _read_text(file, *args, **kwargs):
+            if file == "params.yaml":
+                with original_open(file, *args, **kwargs) as f:
+                    return f.read()
+
+            if not pathlib.Path(file).is_absolute():
+                return self.fs.read_text(file, *args, **kwargs)
+
+            with original_open(file, *args, **kwargs) as f:
+                return f.read()
+
         def _listdir(path, *args, **kwargs):
             if not pathlib.Path(path).is_absolute():
                 return self.fs.listdir(path, detail=False)
@@ -93,8 +106,9 @@ def _listdir(path, *args, **kwargs):
         with unittest.mock.patch("builtins.open", _open):
             with unittest.mock.patch("__main__.open", _open):
                 with unittest.mock.patch("os.listdir", _listdir):
-                    # Jupyter Notebooks replace open with io.open
-                    yield
+                    with unittest.mock.patch("pathlib.Path.read_text", _read_text):
+                        with unittest.mock.patch("pathlib.Path.open", _open):
+                            yield
 
 
 class _NameDescriptor(zninit.Descriptor):