state: load exact module-source overrides before normal sources

pyrefly/lib/state/load.rs

@@ -95,7 +95,15 @@ impl Load {
         timing: Option<&TransactionTimingCounters>,
     ) -> (FileContents, Option<anyhow::Error>) {
         let res = match path.details() {
-            ModulePathDetails::FileSystem(path) => Self::load_from_filesystem(path, timing),
+            ModulePathDetails::FileSystem(path) => {
+                if memory_lookup.filesystem_overlays()
+                    && let Some(contents) = memory_lookup.get(path)
+                {
+                    Ok(contents.as_ref().dupe())
+                } else {
+                    Self::load_from_filesystem(path, timing)
+                }
+            }
             ModulePathDetails::Namespace(_) => Ok(FileContents::from_source("".to_owned())),
             ModulePathDetails::Memory(path) => memory_lookup
                 .get(path)

pyrefly/lib/state/memory.rs

@@ -39,11 +39,25 @@ impl MemoryFilesOverlay {
 pub struct MemoryFilesLookup<'a> {
     base: &'a MemoryFiles,
     overlay: &'a MemoryFilesOverlay,
+    filesystem_overlays: bool,
 }
 
 impl<'a> MemoryFilesLookup<'a> {
     pub fn new(base: &'a MemoryFiles, overlay: &'a MemoryFilesOverlay) -> Self {
-        Self { base, overlay }
+        Self {
+            base,
+            overlay,
+            filesystem_overlays: false,
+        }
+    }
+
+    pub fn with_filesystem_overlays(mut self, enabled: bool) -> Self {
+        self.filesystem_overlays = enabled;
+        self
+    }
+
+    pub fn filesystem_overlays(&self) -> bool {
+        self.filesystem_overlays
     }
 
     pub fn get(&self, path: &Path) -> Option<&'a Arc<FileContents>> {
@@ -53,3 +67,52 @@ impl<'a> MemoryFilesLookup<'a> {
         }
     }
 }
+
+#[cfg(test)]
+mod filesystem_overlay_tests {
+    use std::fs;
+
+    use pyrefly_python::module_path::ModulePath;
+    use tempfile::TempDir;
+
+    use super::*;
+    use crate::state::load::Load;
+
+    #[test]
+    fn filesystem_load_requires_opt_in_and_removal_returns_to_disk() {
+        let directory = TempDir::new().unwrap();
+        let path = directory.path().join("file.py");
+        fs::write(&path, "saved\n").unwrap();
+        let mut base = MemoryFiles::default();
+        let mut initial = MemoryFilesOverlay::default();
+        initial.set(
+            path.clone(),
+            Some(Arc::new(FileContents::from_source(
+                "committed\n".to_owned(),
+            ))),
+        );
+        base.apply_overlay(initial);
+        let mut overlay = MemoryFilesOverlay::default();
+        let module = ModulePath::filesystem(path.clone());
+
+        for (enabled, expected) in [(false, "saved\n"), (true, "committed\n")] {
+            let lookup = MemoryFilesLookup::new(&base, &overlay).with_filesystem_overlays(enabled);
+            let (contents, error) = Load::load_from_path(&module, &lookup, None);
+            assert!(error.is_none(), "{error:?}");
+            assert_eq!(contents, FileContents::from_source(expected.to_owned()));
+        }
+
+        overlay.set(path.clone(), None);
+        let lookup = MemoryFilesLookup::new(&base, &overlay).with_filesystem_overlays(true);
+        let (contents, error) = Load::load_from_path(&module, &lookup, None);
+        assert!(error.is_none(), "{error:?}");
+        assert_eq!(contents, FileContents::from_source("saved\n".to_owned()));
+
+        fs::remove_file(&path).unwrap();
+        let (_, error) = Load::load_from_path(&module, &lookup, None);
+        assert!(
+            error.is_some(),
+            "Missing disk source must not become a successful overlay read"
+        );
+    }
+}

pyrefly/lib/state/state.rs

@@ -46,6 +46,7 @@ use pyrefly_util::demand_tree::DemandEdge;
 use pyrefly_util::demand_tree::DemandSpan;
 use pyrefly_util::events::CategorizedEvents;
 use pyrefly_util::fs_anyhow;
+use pyrefly_util::interned_path::InternedPath;
 use pyrefly_util::lock::Mutex;
 use pyrefly_util::lock::RwLock;
 use pyrefly_util::locked_map::LockedMap;
@@ -2016,6 +2017,7 @@ impl<'a> Transaction<'a> {
 
     fn memory_lookup<'b>(&'b self) -> MemoryFilesLookup<'b> {
         MemoryFilesLookup::new(&self.readable.memory, &self.data.memory_overlay)
+            .with_filesystem_overlays(self.data.state.filesystem_overlays)
     }
 
     fn get_cached_loader(&self, loader: &ArcId<ConfigFile>) -> Arc<LoaderFindCache> {
@@ -2578,22 +2580,31 @@ impl<'a> Transaction<'a> {
         }
     }
 
-    /// Called if the `load_from_memory` portion of loading might have changed.
-    /// Specify which in-memory files might have changed, use None to say they don't exist anymore.
+    /// Update in-memory sources, using the same paths as their module handles.
+    /// None removes a source; with filesystem overlays enabled, the corresponding
+    /// filesystem-backed module returns to disk instead of losing its source.
     pub fn set_memory(&mut self, files: Vec<(PathBuf, Option<Arc<FileContents>>)>) {
         let mut changed = SmallSet::new();
         for (path, contents) in files {
             if self.memory_lookup().get(&path) != contents.as_ref() {
                 self.data.memory_overlay.set(path.clone(), contents);
-                changed.insert(ModulePath::memory(path));
+                changed.insert(InternedPath::new(path));
             }
         }
         self.stats.lock().set_memory_dirty = changed.len();
         if changed.is_empty() {
             return;
         }
+        let filesystem_overlays = self.data.state.filesystem_overlays;
         self.invalidate(
-            |handle| changed.contains(handle.path()),
+            |handle| {
+                let path = match handle.path().details() {
+                    ModulePathDetails::Memory(path) => path,
+                    ModulePathDetails::FileSystem(path) if filesystem_overlays => path,
+                    _ => return false,
+                };
+                changed.contains(path)
+            },
             |state| state.set_dirty_load(),
         );
     }
@@ -3521,6 +3532,7 @@ impl<'a> AsMut<Transaction<'a>> for CancellableTransaction<'a> {
 ///        but they will block committing.
 ///     2. During the committing of `Transaction`, all reads will be blocked.
 pub struct State {
+    filesystem_overlays: bool,
     threads: ThreadPool,
     uniques: UniqueFactory,
     config_finder: ConfigFinder,
@@ -3557,6 +3569,7 @@ impl StateReader<'_> {
 impl State {
     pub fn new(config_finder: ConfigFinder, thread_count: ThreadCount) -> Self {
         Self {
+            filesystem_overlays: false,
             threads: ThreadPool::new(thread_count),
             uniques: UniqueFactory::new(),
             config_finder,
@@ -3567,6 +3580,17 @@ impl State {
         }
     }
 
+    /// Create a state where memory contents override filesystem-backed modules.
+    /// Overlay paths must match resolved file paths; imports still use the filesystem.
+    pub fn new_with_filesystem_overlays(
+        config_finder: ConfigFinder,
+        thread_count: ThreadCount,
+    ) -> Self {
+        let mut state = Self::new(config_finder, thread_count);
+        state.filesystem_overlays = true;
+        state
+    }
+
     pub fn config_finder(&self) -> &ConfigFinder {
         &self.config_finder
     }

pyrefly/lib/test/state.rs

@@ -59,6 +59,183 @@ use crate::state::require::RequireLevels;
 use crate::state::state::State;
 use crate::test::util::TestEnv;
 
+mod filesystem_overlay_tests {
+    use super::*;
+    use crate::state::state::Transaction;
+
+    const SAVED: &str = "value: int = 1\n";
+    const UNSAVED: &str = "value: str = 'unsaved'\n";
+
+    fn setup(root: &Path, overlays: bool) -> (State, [Handle; 2]) {
+        fs::write(root.join("helper.py"), SAVED).unwrap();
+        fs::write(
+            root.join("main.py"),
+            "from helper import value\ncheck: int = value\n",
+        )
+        .unwrap();
+        let mut config = ConfigFile {
+            search_path_from_args: vec![root.to_owned()],
+            disable_search_path_heuristics: true,
+            ..Default::default()
+        };
+        config.python_environment.set_empty_to_default();
+        config.interpreters.skip_interpreter_query = true;
+        assert!(config.configure().is_empty());
+        let sys_info = config.get_sys_info();
+        let finder = ConfigFinder::new_constant(ArcId::new(config));
+        let state = if overlays {
+            State::new_with_filesystem_overlays(finder, TEST_THREAD_COUNT)
+        } else {
+            State::new(finder, TEST_THREAD_COUNT)
+        };
+        let handles = ["main", "helper"].map(|name| {
+            Handle::new(
+                ModuleName::from_str(name),
+                ModulePath::filesystem(root.join(format!("{name}.py"))),
+                sys_info.dupe(),
+            )
+        });
+        (state, handles)
+    }
+
+    fn check(transaction: &mut Transaction, handles: &[Handle; 2], source: &str, errors: usize) {
+        transaction.run(&[handles[0].dupe()], Require::Everything, None);
+        assert_eq!(
+            transaction
+                .get_module_info(&handles[1])
+                .unwrap()
+                .contents()
+                .as_str(),
+            source
+        );
+        assert_eq!(
+            transaction
+                .import_handle(&handles[0], ModuleName::from_str("helper"), None)
+                .finding()
+                .unwrap()
+                .path(),
+            handles[1].path()
+        );
+        let actual = transaction.get_errors([&handles[0]]).collect_errors();
+        assert_eq!(actual.ordinary.len(), errors, "{actual:?}");
+    }
+
+    #[test]
+    fn filesystem_overlays_are_opt_in_and_keep_memory_handles_working() {
+        for enabled in [false, true] {
+            let directory = TempDir::new().unwrap();
+            let (state, handles) = setup(directory.path(), enabled);
+            let path = directory.path().join("helper.py");
+            let memory = Handle::new(
+                ModuleName::from_str("helper"),
+                ModulePath::memory(path.clone()),
+                handles[1].sys_info().dupe(),
+            );
+            let mut transaction = state.new_transaction(Require::Everything, None);
+            transaction.set_memory(vec![(
+                path.clone(),
+                Some(Arc::new(FileContents::from_source(UNSAVED.to_owned()))),
+            )]);
+            check(
+                &mut transaction,
+                &handles,
+                if enabled { UNSAVED } else { SAVED },
+                usize::from(enabled),
+            );
+            transaction.run(&[memory.dupe()], Require::Everything, None);
+            assert_eq!(
+                transaction
+                    .get_module_info(&memory)
+                    .unwrap()
+                    .contents()
+                    .as_str(),
+                UNSAVED
+            );
+            transaction.set_memory(vec![(
+                path.clone(),
+                Some(Arc::new(FileContents::from_source(SAVED.to_owned()))),
+            )]);
+            check(&mut transaction, &handles, SAVED, 0);
+            transaction.run(&[memory.dupe()], Require::Everything, None);
+            assert_eq!(
+                transaction
+                    .get_module_info(&memory)
+                    .unwrap()
+                    .contents()
+                    .as_str(),
+                SAVED
+            );
+            assert_eq!(fs::read_to_string(path).unwrap(), SAVED);
+        }
+    }
+
+    #[test]
+    fn filesystem_overlays_invalidate_importers_and_return_to_disk_on_close() {
+        let directory = TempDir::new().unwrap();
+        let (state, handles) = setup(directory.path(), true);
+        let path = directory.path().join("helper.py");
+        let mut transaction = state.new_committable_transaction(Require::Everything, None);
+        check(transaction.as_mut(), &handles, SAVED, 0);
+        state.commit_transaction(transaction, None);
+
+        let mut transaction = state.new_committable_transaction(Require::Everything, None);
+        transaction.as_mut().set_memory(vec![(
+            path.clone(),
+            Some(Arc::new(FileContents::from_source(UNSAVED.to_owned()))),
+        )]);
+        check(transaction.as_mut(), &handles, UNSAVED, 1);
+        state.commit_transaction(transaction, None);
+
+        const DISK_UPDATE: &str = "value: float = 1.5\n";
+        fs::write(&path, DISK_UPDATE).unwrap();
+        let mut transaction = state.new_committable_transaction(Require::Everything, None);
+        transaction
+            .as_mut()
+            .invalidate_disk(std::slice::from_ref(&path));
+        check(transaction.as_mut(), &handles, UNSAVED, 1);
+        state.commit_transaction(transaction, None);
+
+        let mut transaction = state.new_committable_transaction(Require::Everything, None);
+        transaction.as_mut().set_memory(vec![(path.clone(), None)]);
+        check(transaction.as_mut(), &handles, DISK_UPDATE, 1);
+        state.commit_transaction(transaction, None);
+
+        fs::write(&path, SAVED).unwrap();
+        let mut transaction = state.new_committable_transaction(Require::Everything, None);
+        transaction
+            .as_mut()
+            .invalidate_disk(std::slice::from_ref(&path));
+        check(transaction.as_mut(), &handles, SAVED, 0);
+    }
+
+    #[test]
+    fn filesystem_overlay_pending_removal_shadows_committed_memory_without_committing() {
+        let directory = TempDir::new().unwrap();
+        let (state, handles) = setup(directory.path(), true);
+        let path = directory.path().join("helper.py");
+        let mut transaction = state.new_committable_transaction(Require::Everything, None);
+        transaction.as_mut().set_memory(vec![(
+            path.clone(),
+            Some(Arc::new(FileContents::from_source(UNSAVED.to_owned()))),
+        )]);
+        check(transaction.as_mut(), &handles, UNSAVED, 1);
+        state.commit_transaction(transaction, None);
+
+        {
+            let mut transaction = state.new_transaction(Require::Everything, None);
+            transaction.set_memory(vec![(path.clone(), None)]);
+            check(&mut transaction, &handles, SAVED, 0);
+            transaction.set_memory(vec![(
+                path,
+                Some(Arc::new(FileContents::from_source(SAVED.to_owned()))),
+            )]);
+            check(&mut transaction, &handles, SAVED, 0);
+        }
+        let mut transaction = state.new_transaction(Require::Everything, None);
+        check(&mut transaction, &handles, UNSAVED, 1);
+    }
+}
+
 #[test]
 fn test_fallback_search_path_relative_import_emits_dedicated_diagnostic() {
     let tdir = TempDir::new().unwrap();