fix IDE features should work on unsaved files #1080

lsp/src/extension.ts

@@ -260,6 +260,7 @@ export async function activate(context: ExtensionContext) {
     // Register the server for Python documents
     documentSelector: [
       {scheme: 'file', language: 'python'},
+      {scheme: 'untitled', language: 'python'},
       // Support for notebook cells
       {scheme: 'vscode-notebook-cell', language: 'python'},
     ],

pyrefly/lib/lsp/non_wasm/mod.rs

@@ -12,5 +12,6 @@ pub mod queue;
 pub mod server;
 pub mod stdlib;
 pub mod transaction_manager;
+pub mod unopened_file_tracker;
 pub mod will_rename_files;
 pub mod workspace;

pyrefly/lib/lsp/non_wasm/server.rs

@@ -199,6 +199,7 @@ use crate::lsp::non_wasm::queue::LspQueue;
 use crate::lsp::non_wasm::stdlib::is_python_stdlib_file;
 use crate::lsp::non_wasm::stdlib::should_show_stdlib_error;
 use crate::lsp::non_wasm::transaction_manager::TransactionManager;
+use crate::lsp::non_wasm::unopened_file_tracker::UnopenedFileTracker;
 use crate::lsp::non_wasm::will_rename_files::will_rename_files;
 use crate::lsp::non_wasm::workspace::LspAnalysisConfig;
 use crate::lsp::non_wasm::workspace::Workspace;
@@ -345,6 +346,9 @@ pub struct Server {
     /// should be mapped through here in case they correspond to a cell.
     open_notebook_cells: Arc<RwLock<HashMap<Url, PathBuf>>>,
     open_files: Arc<RwLock<HashMap<PathBuf, Arc<LspFile>>>>,
+    /// Tracks URIs (including virtual/untitled ones) to synthetic on-disk paths so we can
+    /// treat them like regular files throughout the server.
+    unopened_file_tracker: UnopenedFileTracker,
     /// A set of configs where we have already indexed all the files within the config.
     indexed_configs: Mutex<HashSet<ArcId<ConfigFile>>>,
     /// A set of workspaces where we have already performed best-effort indexing.
@@ -609,6 +613,17 @@ pub fn lsp_loop(
 impl Server {
     const FILEWATCHER_ID: &str = "FILEWATCHER";
 
+    fn track_path_for_open(&self, uri: &Url, language_id: &str) -> PathBuf {
+        self.unopened_file_tracker
+            .ensure_path_for_open(uri, language_id)
+    }
+
+    fn require_path_for_uri(&self, uri: &Url) -> anyhow::Result<PathBuf> {
+        self.unopened_file_tracker
+            .path_for_uri(uri)
+            .ok_or_else(|| anyhow::anyhow!("Could not convert uri to filepath: {}", uri))
+    }
+
     fn extract_request_params_or_send_err_response<T>(
         &self,
         params: Result<T::Params, serde_json::Error>,
@@ -667,12 +682,20 @@ impl Server {
                 }
             }
             LspEvent::DidOpenTextDocument(params) => {
-                let contents = Arc::new(LspFile::from_source(params.text_document.text));
+                let lsp_types::DidOpenTextDocumentParams { text_document } = params;
+                let lsp_types::TextDocumentItem {
+                    uri,
+                    language_id,
+                    version,
+                    text,
+                } = text_document;
+                let contents = Arc::new(LspFile::from_source(text));
                 self.did_open(
                     ide_transaction_manager,
                     subsequent_mutation,
-                    params.text_document.uri,
-                    params.text_document.version,
+                    uri,
+                    &language_id,
+                    version,
                     contents,
                 )?;
             }
@@ -700,9 +723,7 @@ impl Server {
                     .collect();
                 let ruff_notebook = params.notebook_document.to_ruff_notebook(&cell_contents)?;
                 let lsp_notebook = LspNotebook::new(ruff_notebook, notebook_document);
-                let notebook_path = url
-                    .to_file_path()
-                    .map_err(|_| anyhow::anyhow!("Could not convert uri to filepath: {}", url))?;
+                let notebook_path = self.track_path_for_open(&url, "python-notebook");
                 for cell_url in lsp_notebook.cell_urls() {
                     self.open_notebook_cells
                         .write()
@@ -712,6 +733,7 @@ impl Server {
                     ide_transaction_manager,
                     subsequent_mutation,
                     url,
+                    "python-notebook",
                     version,
                     Arc::new(LspFile::Notebook(Arc::new(lsp_notebook))),
                 )?;
@@ -1141,7 +1163,7 @@ impl Server {
         {
             folders
                 .iter()
-                .map(|x| x.uri.to_file_path().unwrap())
+                .filter_map(|x| x.uri.to_file_path().ok())
                 .collect()
         } else {
             Vec::new()
@@ -1163,6 +1185,7 @@ impl Server {
             state: Arc::new(State::new(config_finder)),
             open_notebook_cells: Arc::new(RwLock::new(HashMap::new())),
             open_files: Arc::new(RwLock::new(HashMap::new())),
+            unopened_file_tracker: UnopenedFileTracker::new(),
             indexed_configs: Mutex::new(HashSet::new()),
             indexed_workspaces: Mutex::new(HashSet::new()),
             cancellation_handles: Arc::new(Mutex::new(HashMap::new())),
@@ -1195,6 +1218,27 @@ impl Server {
         self.outgoing_requests.lock().insert(id, request);
     }
 
+    fn publish_diagnostics_for_paths(&self, diags: SmallMap<PathBuf, Vec<Diagnostic>>) {
+        let mut entries: Vec<(PathBuf, Url, Vec<Diagnostic>)> = Vec::with_capacity(diags.len());
+        for (path, diagnostics) in diags {
+            if let Some(uri) = self
+                .unopened_file_tracker
+                .uri_for_path(&path)
+                .or_else(|| Url::from_file_path(&path).ok())
+            {
+                entries.push((path, uri, diagnostics));
+            } else {
+                eprintln!("Unable to convert path to uri: {path:?}");
+            }
+        }
+
+        for (path, uri, diagnostics) in entries {
+            let version = self.version_info.lock().get(&path).copied();
+            self.connection
+                .publish_diagnostics_for_uri(uri, diagnostics, version);
+        }
+    }
+
     /// Run the transaction with the in-memory content of open files. Returns the handles of open files when the transaction is done.
     fn validate_in_memory_for_transaction(
         state: &State,
@@ -1593,30 +1637,31 @@ impl Server {
     }
 
     fn did_save(&self, url: Url) {
-        let file = url.to_file_path().unwrap();
-        self.invalidate(move |t| t.invalidate_disk(&[file]));
+        match self.require_path_for_uri(&url) {
+            Ok(file) => self.invalidate(move |t| t.invalidate_disk(&[file])),
+            Err(err) => info!("Skipping didSave for unknown uri {url}: {err}"),
+        }
     }
 
     fn did_open<'a>(
         &'a self,
         ide_transaction_manager: &mut TransactionManager<'a>,
         subsequent_mutation: bool,
         url: Url,
+        language_id: &str,
         version: i32,
         contents: Arc<LspFile>,
     ) -> anyhow::Result<()> {
-        let uri = url
-            .to_file_path()
-            .map_err(|_| anyhow::anyhow!("Could not convert uri to filepath: {}", url))?;
+        let path = self.track_path_for_open(&url, language_id);
         let config_to_populate_files = if self.indexing_mode != IndexingMode::None
-            && let Some(directory) = uri.as_path().parent()
+            && let Some(directory) = path.as_path().parent()
         {
             self.state.config_finder().directory(directory)
         } else {
             None
         };
-        self.version_info.lock().insert(uri.clone(), version);
-        self.open_files.write().insert(uri.clone(), contents);
+        self.version_info.lock().insert(path.clone(), version);
+        self.open_files.write().insert(path.clone(), contents);
         if !subsequent_mutation {
             // In order to improve perceived startup perf, when a file is opened, we run a
             // non-committing transaction that indexes the file with default require level Exports.
@@ -1630,7 +1675,7 @@ impl Server {
             // of a config file, all features become available when background indexing completes.
             info!(
                 "File {} opened, prepare to validate open files.",
-                uri.display()
+                path.display()
             );
             self.validate_in_memory_without_committing(ide_transaction_manager);
         }
@@ -1648,7 +1693,11 @@ impl Server {
         params: DidChangeTextDocumentParams,
     ) -> anyhow::Result<()> {
         let VersionedTextDocumentIdentifier { uri, version } = params.text_document;
-        let file_path = uri.to_file_path().unwrap();
+        let Some(file_path) = self.unopened_file_tracker.path_for_uri(&uri) else {
+            return Err(anyhow::anyhow!(
+                "Received textDocument/didChange for unknown uri: {uri}"
+            ));
+        };
 
         let mut version_info = self.version_info.lock();
         let old_version = version_info.get(&file_path).unwrap_or(&0);
@@ -1683,7 +1732,7 @@ impl Server {
     ) -> anyhow::Result<()> {
         let uri = params.notebook_document.uri.clone();
         let version = params.notebook_document.version;
-        let file_path = uri.to_file_path().unwrap();
+        let file_path = self.require_path_for_uri(&uri)?;
 
         let mut version_info = self.version_info.lock();
         let old_version = version_info.get(&file_path).unwrap_or(&0);
@@ -1849,30 +1898,37 @@ impl Server {
     }
 
     fn did_close(&self, url: Url) {
-        let uri = url.to_file_path().unwrap();
-        self.version_info.lock().remove(&uri);
+        let Ok(path) = self.require_path_for_uri(&url) else {
+            info!("Skipping didClose for unknown uri {url}");
+            return;
+        };
+        self.version_info.lock().remove(&path);
         let open_files = self.open_files.dupe();
-        if let Some(LspFile::Notebook(notebook)) = open_files.write().remove(&uri).as_deref() {
+        if let Some(LspFile::Notebook(notebook)) = open_files.write().remove(&path).as_deref() {
             for cell in notebook.cell_urls() {
                 self.connection
                     .publish_diagnostics_for_uri(cell.clone(), Vec::new(), None);
                 self.open_notebook_cells.write().remove(cell);
             }
         } else {
             self.connection
-                .publish_diagnostics_for_uri(url, Vec::new(), None);
+                .publish_diagnostics_for_uri(url.clone(), Vec::new(), None);
         }
+        self.unopened_file_tracker.forget_uri_path(&url);
         let state = self.state.dupe();
         let lsp_queue = self.lsp_queue.dupe();
         let open_files = self.open_files.dupe();
         let sourcedb_queue = self.sourcedb_queue.dupe();
         let invalidated_configs = self.invalidated_configs.dupe();
+        let path_for_task = path.clone();
         self.recheck_queue.queue_task(Box::new(move || {
             // Clear out the memory associated with this file.
             // Not a race condition because we immediately call validate_in_memory to put back the open files as they are now.
             // Having the extra file hanging around doesn't harm anything, but does use extra memory.
             let mut transaction = state.new_committable_transaction(Require::indexing(), None);
-            transaction.as_mut().set_memory(vec![(uri, None)]);
+            transaction
+                .as_mut()
+                .set_memory(vec![(path_for_task.clone(), None)]);
             let _ =
                 Self::validate_in_memory_for_transaction(&state, &open_files, transaction.as_mut());
             state.commit_transaction(transaction);
@@ -1945,12 +2001,17 @@ impl Server {
         uri: &Url,
         method: Option<&str>,
     ) -> Option<(Handle, Option<LspAnalysisConfig>)> {
-        let path = self
-            .open_notebook_cells
-            .read()
-            .get(uri)
-            .cloned()
-            .unwrap_or_else(|| uri.to_file_path().unwrap());
+        let path = if let Some(notebook_path) = self.open_notebook_cells.read().get(uri) {
+            notebook_path.clone()
+        } else {
+            match self.require_path_for_uri(uri) {
+                Ok(path) => path,
+                Err(err) => {
+                    info!("Skipping request - invalid uri {uri}: {err}");
+                    return None;
+                }
+            }
+        };
         self.workspaces.get_with(path.clone(), |(_, workspace)| {
             // Check if all language services are disabled
             if workspace.disable_language_services {
@@ -2410,11 +2471,16 @@ impl Server {
             // TODO(yangdanny) handle notebooks
             return None;
         }
+        let path = match self.require_path_for_uri(uri) {
+            Ok(path) => path,
+            Err(err) => {
+                info!("Skipping document symbols for unknown uri {uri}: {err}");
+                return None;
+            }
+        };
         if self
             .workspaces
-            .get_with(uri.to_file_path().unwrap(), |(_, workspace)| {
-                workspace.disable_language_services
-            })
+            .get_with(path, |(_, workspace)| workspace.disable_language_services)
             || !self
                 .initialize_params
                 .capabilities
@@ -2593,7 +2659,19 @@ impl Server {
             cell_uri = Some(uri);
             notebook_path.as_path().to_owned()
         } else {
-            uri.to_file_path().unwrap()
+            match self.require_path_for_uri(uri) {
+                Ok(path) => path,
+                Err(err) => {
+                    info!("Skipping diagnostics for unknown uri {uri}: {err}");
+                    return DocumentDiagnosticReport::Full(RelatedFullDocumentDiagnosticReport {
+                        full_document_diagnostic_report: FullDocumentDiagnosticReport {
+                            items: Vec::new(),
+                            result_id: None,
+                        },
+                        related_documents: None,
+                    });
+                }
+            }
         };
         let handle = make_open_handle(&self.state, &path);
         let mut items = Vec::new();