From b4c42eb83b1215ad980a11265f38d84bdeba7e28 Mon Sep 17 00:00:00 2001 From: Dhruv Manilawala Date: Mon, 14 Jul 2025 20:17:51 +0530 Subject: [PATCH] [ty] Add virtual files to the only project database (#19322) ## Summary Previously, the virtual files were being added to the default database that's present on the session. This is wrong because the default database is for any files that don't belong to any project i.e., they're outside of any projects managed by the server. Virtual files are neither part of the project nor it is outside the projects. This was not the intention as in the initial version, virtual files were being added to the only project database managed by the server. This PR fixes this by reverting back to the original behavior where virtual files will be added to the only project database present. When support for multiple workspace and project is added, this will require updating (https://github.com/astral-sh/ty/issues/794). This is required for #19264 because workspace diagnostics doesn't check the default project database yet. Ideally, the default db should be checked as well. The implementation of this PR means that virtual files are now being included for workspace diagnostics but it doesn't work completely e.g., if I save an untitled file the diagnostics disappears but it doesn't appear back for the (now) saved file on disk as shown in the following video demonstration: https://github.com/user-attachments/assets/123e8d20-1e95-4c7d-b7eb-eb65be8c476e --- crates/ty_server/src/server/api.rs | 9 +- .../ty_server/src/server/api/diagnostics.rs | 2 +- .../server/api/notifications/did_change.rs | 4 +- .../src/server/api/notifications/did_open.rs | 8 +- .../api/notifications/did_open_notebook.rs | 2 +- crates/ty_server/src/session.rs | 82 ++++++++++--------- 6 files changed, 52 insertions(+), 55 deletions(-) diff --git a/crates/ty_server/src/server/api.rs b/crates/ty_server/src/server/api.rs index a3a55c25e2..b9dd261223 100644 --- a/crates/ty_server/src/server/api.rs +++ b/crates/ty_server/src/server/api.rs @@ -236,14 +236,7 @@ where }); }; - let db = match &path { - AnySystemPath::System(path) => match session.project_db_for_path(path) { - Some(db) => db.clone(), - None => session.default_project_db().clone(), - }, - AnySystemPath::SystemVirtual(_) => session.default_project_db().clone(), - }; - + let db = session.project_db(&path).clone(); let snapshot = session.take_document_snapshot(url); Box::new(move |client| { diff --git a/crates/ty_server/src/server/api/diagnostics.rs b/crates/ty_server/src/server/api/diagnostics.rs index f7cd24bc7e..fe585ff105 100644 --- a/crates/ty_server/src/server/api/diagnostics.rs +++ b/crates/ty_server/src/server/api/diagnostics.rs @@ -82,7 +82,7 @@ pub(super) fn publish_diagnostics(session: &Session, key: &DocumentKey, client: } }; - let db = session.project_db_or_default(key.path()); + let db = session.project_db(key.path()); let Some(diagnostics) = compute_diagnostics(db, &snapshot) else { return; diff --git a/crates/ty_server/src/server/api/notifications/did_change.rs b/crates/ty_server/src/server/api/notifications/did_change.rs index e9042da999..68f6f883e0 100644 --- a/crates/ty_server/src/server/api/notifications/did_change.rs +++ b/crates/ty_server/src/server/api/notifications/did_change.rs @@ -41,8 +41,8 @@ impl SyncNotificationHandler for DidChangeTextDocumentHandler { .with_failure_code(ErrorCode::InternalError)?; let changes = match key.path() { - AnySystemPath::System(path) => { - vec![ChangeEvent::file_content_changed(path.clone())] + AnySystemPath::System(system_path) => { + vec![ChangeEvent::file_content_changed(system_path.clone())] } AnySystemPath::SystemVirtual(virtual_path) => { vec![ChangeEvent::ChangedVirtual(virtual_path.clone())] diff --git a/crates/ty_server/src/server/api/notifications/did_open.rs b/crates/ty_server/src/server/api/notifications/did_open.rs index 6ef920b535..dfb1e1b472 100644 --- a/crates/ty_server/src/server/api/notifications/did_open.rs +++ b/crates/ty_server/src/server/api/notifications/did_open.rs @@ -44,12 +44,14 @@ impl SyncNotificationHandler for DidOpenTextDocumentHandler { let document = TextDocument::new(text, version).with_language_id(&language_id); session.open_text_document(key.path(), document); - match key.path() { + let path = key.path(); + + match path { AnySystemPath::System(system_path) => { - session.apply_changes(key.path(), vec![ChangeEvent::Opened(system_path.clone())]); + session.apply_changes(path, vec![ChangeEvent::Opened(system_path.clone())]); } AnySystemPath::SystemVirtual(virtual_path) => { - let db = session.default_project_db_mut(); + let db = session.project_db_mut(path); db.files().virtual_file(db, virtual_path); } } diff --git a/crates/ty_server/src/server/api/notifications/did_open_notebook.rs b/crates/ty_server/src/server/api/notifications/did_open_notebook.rs index c3e4fff5cd..201add9587 100644 --- a/crates/ty_server/src/server/api/notifications/did_open_notebook.rs +++ b/crates/ty_server/src/server/api/notifications/did_open_notebook.rs @@ -43,7 +43,7 @@ impl SyncNotificationHandler for DidOpenNotebookHandler { session.apply_changes(&path, vec![ChangeEvent::Opened(system_path.clone())]); } AnySystemPath::SystemVirtual(virtual_path) => { - let db = session.default_project_db_mut(); + let db = session.project_db_mut(&path); db.files().virtual_file(db, virtual_path); } } diff --git a/crates/ty_server/src/session.rs b/crates/ty_server/src/session.rs index c31ec23884..c71e1e7301 100644 --- a/crates/ty_server/src/session.rs +++ b/crates/ty_server/src/session.rs @@ -163,16 +163,47 @@ impl Session { &self.workspaces } - // TODO(dhruvmanila): Ideally, we should have a single method for `workspace_db_for_path_mut` - // and `default_workspace_db_mut` but the borrow checker doesn't allow that. - // https://github.com/astral-sh/ruff/pull/13041#discussion_r1726725437 + /// Returns a reference to the project's [`ProjectDatabase`] in which the given `path` belongs. + /// + /// If the path is a system path, it will return the project database that is closest to the + /// given path, or the default project if no project is found for the path. + /// + /// If the path is a virtual path, it will return the first project database in the session. + pub(crate) fn project_db(&self, path: &AnySystemPath) -> &ProjectDatabase { + match path { + AnySystemPath::System(system_path) => self + .project_db_for_path(system_path) + .unwrap_or_else(|| self.default_project.get(self.index.as_ref())), + AnySystemPath::SystemVirtual(_virtual_path) => { + // TODO: Currently, ty only supports single workspace but we need to figure out + // which project should this virtual path belong to when there are multiple + // projects: https://github.com/astral-sh/ty/issues/794 + self.projects.iter().next().map(|(_, db)| db).unwrap() + } + } + } - /// Returns a reference to the project's [`ProjectDatabase`] corresponding to the given path, - /// or the default project if no project is found for the path. - pub(crate) fn project_db_or_default(&self, path: &AnySystemPath) -> &ProjectDatabase { - path.as_system() - .and_then(|path| self.project_db_for_path(path)) - .unwrap_or_else(|| self.default_project_db()) + /// Returns a mutable reference to the project's [`ProjectDatabase`] in which the given `path` + /// belongs. + /// + /// Refer to [`project_db`] for more details on how the project is selected. + /// + /// [`project_db`]: Session::project_db + pub(crate) fn project_db_mut(&mut self, path: &AnySystemPath) -> &mut ProjectDatabase { + match path { + AnySystemPath::System(system_path) => self + .projects + .range_mut(..=system_path.to_path_buf()) + .next_back() + .map(|(_, db)| db) + .unwrap_or_else(|| self.default_project.get_mut(self.index.as_ref())), + AnySystemPath::SystemVirtual(_virtual_path) => { + // TODO: Currently, ty only supports single workspace but we need to figure out + // which project should this virtual path belong to when there are multiple + // projects: https://github.com/astral-sh/ty/issues/794 + self.projects.iter_mut().next().map(|(_, db)| db).unwrap() + } + } } /// Returns a reference to the project's [`ProjectDatabase`] corresponding to the given path, if @@ -187,18 +218,6 @@ impl Session { .map(|(_, db)| db) } - /// Returns a mutable reference to the project [`ProjectDatabase`] corresponding to the given - /// path, if any. - pub(crate) fn project_db_for_path_mut( - &mut self, - path: impl AsRef, - ) -> Option<&mut ProjectDatabase> { - self.projects - .range_mut(..=path.as_ref().to_path_buf()) - .next_back() - .map(|(_, db)| db) - } - pub(crate) fn apply_changes( &mut self, path: &AnySystemPath, @@ -212,25 +231,8 @@ impl Session { .cloned() }); - let db = match path { - AnySystemPath::System(path) => match self.project_db_for_path_mut(path) { - Some(db) => db, - None => self.default_project_db_mut(), - }, - AnySystemPath::SystemVirtual(_) => self.default_project_db_mut(), - }; - - db.apply_changes(changes, overrides.as_ref()) - } - - /// Returns a reference to the default project [`ProjectDatabase`]. - pub(crate) fn default_project_db(&self) -> &ProjectDatabase { - self.default_project.get(self.index.as_ref()) - } - - /// Returns a mutable reference to the default project [`ProjectDatabase`]. - pub(crate) fn default_project_db_mut(&mut self) -> &mut ProjectDatabase { - self.default_project.get_mut(self.index.as_ref()) + self.project_db_mut(path) + .apply_changes(changes, overrides.as_ref()) } /// Returns a mutable iterator over all project databases that have been initialized to this point.