mount: keep custom CWD on SW reset during session when is mounted and CWD has been modified with command

Author: pavelrevak

README_bench.md

@@ -85,6 +85,13 @@ Reading files from mounted VFS (PC filesystem accessible on device).
 | Read 50 x 2KB files | **0.6s** | 3.2s | **5.6x** |
 | Read 1 x 100KB file | **0.3s** | 0.5s | **1.8x** |
 
+### ESP32-S3 - USB-CDC - MacOS
+
+| Test | mpytool | mpremote | Speedup |
+|------|---------|----------|---------|
+| Read 50 x 2KB files | **1.6s** | 3.9s | **2.3x** |
+| Read 1 x 100KB file | **0.6s** | 1.0s | **1.7x** |
+
 ### Summary
 
 Batch LISTDIR (1 RTT) vs iterative ilistdir (N+1 RTT) makes significant difference when opening many files. The speedup is more pronounced on faster CPUs (ESP32-C6 160MHz vs RP2040 133MHz).

mpytool/mpy.py

@@ -144,6 +144,8 @@ def __init__(self, conn, log=None, chunk_size=None):
         # Multi-mount state
         self._mounts = []  # [(mid, mount_point, local_path, handler), ...]
         self._next_mid = 0
+        self._custom_cwd = None  # CWD to restore after soft reset
+        self._custom_syspath = None  # sys.path to restore after soft reset
         # VfsProtocol always registered - handles stale VFS and all mounts
         self._vfs_protocol = _mount.VfsProtocol(
             conn, remount_fn=self._do_remount_all, log=log)
@@ -328,6 +330,7 @@ def chdir(self, path):
         self.import_module('os')
         try:
             self._mpy_comm.exec(f"os.chdir('{_escape_path(path)}')")
+            self._custom_cwd = path
         except _mpy_comm.CmdError as err:
             raise DirNotFound(path) from err
 
@@ -347,7 +350,9 @@ def set_sys_path(self, *paths):
             *paths: path strings
         """
         self.import_module('sys')
-        self._mpy_comm.exec(f"sys.path = {list(paths)!r}")
+        new_path = list(paths)
+        self._mpy_comm.exec(f"sys.path = {new_path!r}")
+        self._custom_syspath = new_path
 
     def prepend_sys_path(self, *paths):
         """Add paths to beginning of sys.path (removes duplicates)
@@ -357,7 +362,9 @@ def prepend_sys_path(self, *paths):
         """
         current = self.get_sys_path()
         current = [p for p in current if p not in paths]
-        self._mpy_comm.exec(f"sys.path = {list(paths) + current!r}")
+        new_path = list(paths) + current
+        self._mpy_comm.exec(f"sys.path = {new_path!r}")
+        self._custom_syspath = new_path
 
     def append_sys_path(self, *paths):
         """Add paths to end of sys.path (removes duplicates)
@@ -367,7 +374,9 @@ def append_sys_path(self, *paths):
         """
         current = self.get_sys_path()
         current = [p for p in current if p not in paths]
-        self._mpy_comm.exec(f"sys.path = {current + list(paths)!r}")
+        new_path = current + list(paths)
+        self._mpy_comm.exec(f"sys.path = {new_path!r}")
+        self._custom_syspath = new_path
 
     def remove_from_sys_path(self, *paths):
         """Remove specified paths from sys.path
@@ -378,6 +387,7 @@ def remove_from_sys_path(self, *paths):
         current = self.get_sys_path()
         new_path = [p for p in current if p not in paths]
         self._mpy_comm.exec(f"sys.path = {new_path!r}")
+        self._custom_syspath = new_path
 
     def hashfile(self, path):
         """Compute SHA256 hash of file
@@ -1576,6 +1586,12 @@ def mount(
 
         self._mounts.append((mid, mount_point, local_path, handler))
         self._next_mid += 1
+
+        # First mount: change CWD to mount point (mpremote compatibility)
+        if len(self._mounts) == 1:
+            self._mpy_comm.exec(f"os.chdir('{mp_escaped}')", timeout=3)
+            self._custom_cwd = mount_point
+
         return handler
 
     def _do_remount_all(self):
@@ -1595,12 +1611,15 @@ def _do_remount_all(self):
             self._mpy_comm.exec(
                 f"_mt_mount('{mp_escaped}',{mid},{self._chunk_size})",
                 timeout=5)
-        # Restore CWD to first mount point (mpremote compatibility)
-        if self._mounts:
-            first_mount = self._mounts[0]
-            if first_mount[0] is not None:  # Not a submount
-                mp_escaped = _escape_path(first_mount[1])
-                self._mpy_comm.exec(
-                    f"os.chdir('{mp_escaped}')", timeout=3)
+        # Restore CWD
+        if self._custom_cwd is not None:
+            cwd_escaped = _escape_path(self._custom_cwd)
+            self._mpy_comm.exec(f"os.chdir('{cwd_escaped}')", timeout=3)
+
+        # Restore sys.path
+        if self._custom_syspath is not None:
+            self._mpy_comm.exec(
+                f"import sys; sys.path[:] = {self._custom_syspath!r}", timeout=3)
+
         self._mpy_comm.exit_raw_repl()
 

tests/test_mount.py

@@ -1960,5 +1960,103 @@ def test_readline_after_seek(self):
         self.assertEqual(result[4:4+length], b'BBB\n')
 
 
+class TestMountCwdTracking(unittest.TestCase):
+    """Tests for mount() CWD tracking"""
+
+    def setUp(self):
+        self.mock_conn = Mock()
+        self.mock_conn._escape_handlers = {}
+        self.mock_conn.register_escape_handler = Mock(
+            side_effect=lambda b, h: self.mock_conn._escape_handlers.__setitem__(b, h))
+        from mpytool.mpy import Mpy
+        self.mpy = Mpy(self.mock_conn)
+        self.mpy._mpy_comm = Mock()
+        self.mpy._mpy_comm.exec_eval.return_value = 4096  # chunk size
+        self.temp_dir = tempfile.mkdtemp()
+
+    def tearDown(self):
+        shutil.rmtree(self.temp_dir)
+
+    def test_first_mount_saves_cwd(self):
+        """First mount saves mount_point to _custom_cwd"""
+        self.mpy.mount(self.temp_dir, '/remote')
+        self.assertEqual(self.mpy._custom_cwd, '/remote')
+
+    def test_second_mount_does_not_change_cwd(self):
+        """Second mount does not overwrite _custom_cwd"""
+        temp_dir2 = tempfile.mkdtemp()
+        try:
+            self.mpy.mount(self.temp_dir, '/remote')
+            self.mpy.mount(temp_dir2, '/other')
+            # CWD should still be from first mount
+            self.assertEqual(self.mpy._custom_cwd, '/remote')
+        finally:
+            shutil.rmtree(temp_dir2)
+
+
+class TestRemountRestoreState(unittest.TestCase):
+    """Tests for _do_remount_all() state restore"""
+
+    def setUp(self):
+        self.mock_conn = Mock()
+        self.mock_conn._escape_handlers = {}
+        self.mock_conn.register_escape_handler = Mock(
+            side_effect=lambda b, h: self.mock_conn._escape_handlers.__setitem__(b, h))
+        from mpytool.mpy import Mpy
+        self.mpy = Mpy(self.mock_conn)
+        self.mpy._mpy_comm = Mock()
+        self.mpy._mpy_comm.exec_eval.return_value = 4096
+        self.temp_dir = tempfile.mkdtemp()
+
+    def tearDown(self):
+        shutil.rmtree(self.temp_dir)
+
+    def test_remount_restores_cwd(self):
+        """_do_remount_all() restores _custom_cwd"""
+        self.mpy.mount(self.temp_dir, '/remote')
+        self.mpy._custom_cwd = '/lib'  # user changed CWD
+        self.mpy._mpy_comm.reset_mock()
+
+        self.mpy._do_remount_all()
+
+        # Check that os.chdir('/lib') was called
+        exec_calls = [str(c) for c in self.mpy._mpy_comm.exec.call_args_list]
+        self.assertTrue(
+            any("os.chdir('/lib')" in c for c in exec_calls),
+            f"Expected os.chdir('/lib') in {exec_calls}")
+
+    def test_remount_restores_syspath(self):
+        """_do_remount_all() restores _custom_syspath"""
+        self.mpy.mount(self.temp_dir, '/remote')
+        self.mpy._custom_syspath = ['/lib', '/app']
+        self.mpy._mpy_comm.reset_mock()
+
+        self.mpy._do_remount_all()
+
+        # Check that sys.path was restored
+        exec_calls = [str(c) for c in self.mpy._mpy_comm.exec.call_args_list]
+        self.assertTrue(
+            any("sys.path[:] = ['/lib', '/app']" in c for c in exec_calls),
+            f"Expected sys.path restore in {exec_calls}")
+
+    def test_remount_no_restore_if_none(self):
+        """_do_remount_all() skips restore if state is None"""
+        self.mpy.mount(self.temp_dir, '/remote')
+        self.mpy._custom_cwd = None
+        self.mpy._custom_syspath = None
+        self.mpy._mpy_comm.reset_mock()
+
+        self.mpy._do_remount_all()
+
+        # Check that os.chdir and sys.path were NOT called
+        exec_calls = [str(c) for c in self.mpy._mpy_comm.exec.call_args_list]
+        self.assertFalse(
+            any("os.chdir" in c for c in exec_calls),
+            f"os.chdir should not be called when _custom_cwd is None")
+        self.assertFalse(
+            any("sys.path[:] =" in c for c in exec_calls),
+            f"sys.path should not be restored when _custom_syspath is None")
+
+
 if __name__ == '__main__':
     unittest.main()

tests/test_mpy.py

@@ -604,5 +604,60 @@ def test_try_raw_paste_timeout_zero_fallback_exec(self):
         self.assertFalse(self.comm._repl_mode)
 
 
+class TestCwdAndPathTracking(unittest.TestCase):
+    """Tests for CWD and sys.path tracking for soft reset restore"""
+
+    def setUp(self):
+        self.mock_conn = Mock()
+        self.mock_conn._escape_handlers = {}
+        self.mock_conn.register_escape_handler = Mock(
+            side_effect=lambda b, h: self.mock_conn._escape_handlers.__setitem__(b, h))
+        self.mpy = Mpy(self.mock_conn)
+
+    def test_initial_state_none(self):
+        """_custom_cwd and _custom_syspath are None initially"""
+        self.assertIsNone(self.mpy._custom_cwd)
+        self.assertIsNone(self.mpy._custom_syspath)
+
+    def test_chdir_saves_cwd(self):
+        """chdir() saves path to _custom_cwd"""
+        self.mpy._mpy_comm = Mock()
+        self.mpy._imported = ['os']
+        self.mpy.chdir('/lib')
+        self.assertEqual(self.mpy._custom_cwd, '/lib')
+        self.mpy._mpy_comm.exec.assert_called()
+
+    def test_set_sys_path_saves_path(self):
+        """set_sys_path() saves to _custom_syspath"""
+        self.mpy._mpy_comm = Mock()
+        self.mpy._imported = ['sys']
+        self.mpy.set_sys_path('/lib', '/app')
+        self.assertEqual(self.mpy._custom_syspath, ['/lib', '/app'])
+
+    def test_prepend_sys_path_saves_path(self):
+        """prepend_sys_path() saves combined path to _custom_syspath"""
+        self.mpy._mpy_comm = Mock()
+        self.mpy._mpy_comm.exec_eval.return_value = ['/old']
+        self.mpy._imported = ['sys']
+        self.mpy.prepend_sys_path('/new')
+        self.assertEqual(self.mpy._custom_syspath, ['/new', '/old'])
+
+    def test_append_sys_path_saves_path(self):
+        """append_sys_path() saves combined path to _custom_syspath"""
+        self.mpy._mpy_comm = Mock()
+        self.mpy._mpy_comm.exec_eval.return_value = ['/old']
+        self.mpy._imported = ['sys']
+        self.mpy.append_sys_path('/new')
+        self.assertEqual(self.mpy._custom_syspath, ['/old', '/new'])
+
+    def test_remove_from_sys_path_saves_path(self):
+        """remove_from_sys_path() saves filtered path to _custom_syspath"""
+        self.mpy._mpy_comm = Mock()
+        self.mpy._mpy_comm.exec_eval.return_value = ['/a', '/b', '/c']
+        self.mpy._imported = ['sys']
+        self.mpy.remove_from_sys_path('/b')
+        self.assertEqual(self.mpy._custom_syspath, ['/a', '/c'])
+
+
 if __name__ == "__main__":
     unittest.main()