Fix MemoryFileSystem._internal_path prefix stripping corruption

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

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')