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.