[ty] Add virtual files to the only project database

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| {

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;

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())]

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);
             }
         }

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);
             }
         }

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<SystemPath>,
-    ) -> 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.