diff --git a/pyglove/core/io/file_system.py b/pyglove/core/io/file_system.py index 21e07ce..4730f7d 100644 --- a/pyglove/core/io/file_system.py +++ b/pyglove/core/io/file_system.py @@ -311,7 +311,16 @@ def __init__(self, prefix: str = '/mem/'): self._prefix = prefix def _internal_path(self, path: Union[str, os.PathLike[str]]) -> str: - return '/' + resolve_path(path).lstrip(self._prefix) + path = resolve_path(path) + # ``str.lstrip`` treats its argument as a set of characters, so it must not + # be used here: it would incorrectly strip path components that begin with + # the same characters as the prefix (e.g. ``/mem/models`` -> ``odels``). + if path.startswith(self._prefix): + path = path[len(self._prefix):] + elif path == self._prefix.rstrip('/'): + # ``/mem`` (the prefix without the trailing slash) refers to the root. + path = '' + return '/' + path def _locate(self, path: Union[str, os.PathLike[str]]) -> Any: current = self._root diff --git a/pyglove/core/io/file_system_test.py b/pyglove/core/io/file_system_test.py index 5ea8d06..240df4a 100644 --- a/pyglove/core/io/file_system_test.py +++ b/pyglove/core/io/file_system_test.py @@ -290,6 +290,23 @@ def test_file_system(self): fs.rmdirs(os.path.join(dir_a, 'b/c')) self.assertEqual(fs.listdir(dir_a), ['file1']) # pylint: disable=g-generic-assert + def test_internal_path_prefix_stripping(self): + fs = file_system.MemoryFileSystem() + # Path components beginning with the same characters as the ``/mem/`` + # prefix (e.g. ``m`` or ``e``) must not be stripped away. + fs.mkdirs('/mem/models/a') + fs.mkdirs('/mem/eval/b') + self.assertEqual(sorted(fs.listdir('/mem')), ['eval', 'models']) + self.assertTrue(fs.exists('/mem/models/a')) + self.assertTrue(fs.exists('/mem/eval/b')) + + # Paths whose first component starts with prefix characters must not + # collide with shorter paths. + fs.mkdirs('/mem/odels/c') + self.assertTrue(fs.exists('/mem/odels/c')) + self.assertEqual( + sorted(fs.listdir('/mem')), ['eval', 'models', 'odels']) + def test_glob(self): fs = file_system.MemoryFileSystem() fs.mkdirs('/mem/a/b/c')