[ty] Fix handling of non-Python text documents

crates/ty_server/src/document/text_document.rs

@@ -24,7 +24,7 @@ pub struct TextDocument {
     version: DocumentVersion,
 
     /// The language ID of the document as provided by the client.
-    language_id: Option<LanguageId>,
+    language_id: LanguageId,
 
     /// For cells, the path to the notebook document.
     notebook: Option<AnySystemPath>,
@@ -46,22 +46,16 @@ impl From<&str> for LanguageId {
 }
 
 impl TextDocument {
-    pub fn new(url: Url, contents: String, version: DocumentVersion) -> Self {
+    pub fn new(url: Url, contents: String, version: DocumentVersion, language_id: &str) -> Self {
         Self {
             url,
             contents,
             version,
-            language_id: None,
+            language_id: LanguageId::from(language_id),
             notebook: None,
         }
     }
 
-    #[must_use]
-    pub fn with_language_id(mut self, language_id: &str) -> Self {
-        self.language_id = Some(LanguageId::from(language_id));
-        self
-    }
-
     #[must_use]
     pub(crate) fn with_notebook(mut self, notebook: AnySystemPath) -> Self {
         self.notebook = Some(notebook);
@@ -84,7 +78,7 @@ impl TextDocument {
         self.version
     }
 
-    pub fn language_id(&self) -> Option<LanguageId> {
+    pub fn language_id(&self) -> LanguageId {
         self.language_id
     }
 
@@ -177,6 +171,7 @@ def interface():
 "#
             .to_string(),
             0,
+            "python",
         );
 
         // Add an `s`, remove it again (back to the original code), and then re-add the `s`

crates/ty_server/src/server/api/notifications/did_open.rs

@@ -2,7 +2,6 @@ use lsp_types::notification::DidOpenTextDocument;
 use lsp_types::{DidOpenTextDocumentParams, TextDocumentItem};
 
 use crate::TextDocument;
-use crate::document::LanguageId;
 use crate::server::Result;
 use crate::server::api::diagnostics::publish_diagnostics_if_needed;
 use crate::server::api::traits::{NotificationHandler, SyncNotificationHandler};
@@ -31,11 +30,7 @@ impl SyncNotificationHandler for DidOpenTextDocumentHandler {
                 },
         } = params;
 
-        let text_doc = TextDocument::new(uri, text, version).with_language_id(&language_id);
-        if matches!(text_doc.language_id(), Some(LanguageId::Other)) {
-            return Ok(());
-        }
-
+        let text_doc = TextDocument::new(uri, text, version, &language_id);
         let document = session.open_text_document(text_doc);
         publish_diagnostics_if_needed(&document, session, client);
 

crates/ty_server/src/server/api/notifications/did_open_notebook.rs

@@ -39,9 +39,9 @@ impl SyncNotificationHandler for DidOpenNotebookHandler {
         let notebook_path = document.notebook_or_file_path();
 
         for cell in params.cell_text_documents {
-            let cell_document = TextDocument::new(cell.uri, cell.text, cell.version)
-                .with_language_id(&cell.language_id)
-                .with_notebook(notebook_path.clone());
+            let cell_document =
+                TextDocument::new(cell.uri, cell.text, cell.version, &cell.language_id)
+                    .with_notebook(notebook_path.clone());
             session.open_text_document(cell_document);
         }
 

crates/ty_server/src/session.rs

@@ -33,7 +33,7 @@ pub(crate) use self::options::InitializationOptions;
 pub use self::options::{ClientOptions, DiagnosticMode, GlobalOptions, WorkspaceOptions};
 pub(crate) use self::settings::{GlobalSettings, WorkspaceSettings};
 use crate::capabilities::{ResolvedClientCapabilities, server_diagnostic_options};
-use crate::document::{DocumentKey, DocumentVersion, NotebookDocument};
+use crate::document::{DocumentKey, DocumentVersion, LanguageId, NotebookDocument};
 use crate::server::{Action, publish_settings_diagnostics};
 use crate::session::client::Client;
 use crate::session::index::Document;
@@ -1166,7 +1166,7 @@ impl Session {
     /// Returns a handle to the opened document.
     pub(crate) fn open_notebook_document(&mut self, document: NotebookDocument) -> DocumentHandle {
         let handle = self.index_mut().open_notebook_document(document);
-        self.open_document_in_db(&handle);
+        self.open_document_in_db(&handle, None);
         handle
     }
 
@@ -1175,12 +1175,13 @@ impl Session {
     ///
     /// Returns a handle to the opened document.
     pub(crate) fn open_text_document(&mut self, document: TextDocument) -> DocumentHandle {
+        let language_id = document.language_id();
         let handle = self.index_mut().open_text_document(document);
-        self.open_document_in_db(&handle);
+        self.open_document_in_db(&handle, Some(language_id));
         handle
     }
 
-    fn open_document_in_db(&mut self, document: &DocumentHandle) {
+    fn open_document_in_db(&mut self, document: &DocumentHandle, language_id: Option<LanguageId>) {
         let path = document.notebook_or_file_path();
 
         // This is a "maybe" because the `File` might've not been interned yet i.e., the
@@ -1193,6 +1194,11 @@ impl Session {
                 .is_none_or(|file| !file.exists(db))
         });
 
+        // When we know the document isn't a Python source file
+        // then we'll avoid adding it to the project. (But we
+        // still track it as part of the index.)
+        let is_not_python = matches!(language_id, Some(LanguageId::Other));
+
         match path {
             AnySystemPath::System(system_path) => {
                 let event = if is_maybe_new_system_file {
@@ -1205,6 +1211,10 @@ impl Session {
                 };
                 self.apply_changes(path, vec![event]);
 
+                if is_not_python {
+                    return;
+                }
+
                 let db = self.project_db_mut(path);
                 match system_path_to_file(db, system_path) {
                     Ok(file) => {
@@ -1220,6 +1230,10 @@ impl Session {
                 }
             }
             AnySystemPath::SystemVirtual(virtual_path) => {
+                if is_not_python {
+                    return;
+                }
+
                 let db = self.project_db_mut(path);
                 let virtual_file = db.files().virtual_file(db, virtual_path);
                 db.project().open_file(db, virtual_file.file());

crates/ty_server/src/session/index.rs

@@ -110,10 +110,14 @@ impl Index {
             self.documents.insert(
                 DocumentKey::from_url(&opened_cell.uri),
                 Document::Text(
-                    TextDocument::new(opened_cell.uri, opened_cell.text, opened_cell.version)
-                        .with_language_id(&opened_cell.language_id)
-                        .with_notebook(notebook_path.clone())
-                        .into(),
+                    TextDocument::new(
+                        opened_cell.uri,
+                        opened_cell.text,
+                        opened_cell.version,
+                        &opened_cell.language_id,
+                    )
+                    .with_notebook(notebook_path.clone())
+                    .into(),
                 ),
             );
         }
@@ -231,7 +235,7 @@ impl Document {
 
     pub(crate) fn language_id(&self) -> Option<LanguageId> {
         match self {
-            Self::Text(document) => document.language_id(),
+            Self::Text(document) => Some(document.language_id()),
             Self::Notebook(_) => None,
         }
     }

crates/ty_server/src/system.rs

@@ -124,7 +124,7 @@ impl LSPSystem {
         extension: Option<&str>,
     ) -> Option<PySourceType> {
         match document {
-            Document::Text(text) => match text.language_id()? {
+            Document::Text(text) => match text.language_id() {
                 LanguageId::Python => Some(
                     extension
                         .and_then(PySourceType::try_from_extension)

crates/ty_server/tests/e2e/publish_diagnostics.rs

@@ -429,6 +429,8 @@ def foo() -> str:
     insta::assert_debug_snapshot!(diagnostics);
 
     server.close_text_document(foo);
+    let diagnostics = server.await_notification::<PublishDiagnostics>();
+    insta::assert_debug_snapshot!(diagnostics);
 
     let params = DidOpenTextDocumentParams {
         text_document: TextDocumentItem {

crates/ty_server/tests/e2e/snapshots/e2e__publish_diagnostics__changing_language_of_file_without_extension-3.snap

@@ -0,0 +1,21 @@
+---
+source: crates/ty_server/tests/e2e/publish_diagnostics.rs
+expression: diagnostics
+---
+PublishDiagnosticsParams {
+    uri: Url {
+        scheme: "file",
+        cannot_be_a_base: false,
+        username: "",
+        password: None,
+        host: None,
+        port: None,
+        path: "<temp_dir>/src/foo",
+        query: None,
+        fragment: None,
+    },
+    diagnostics: [],
+    version: Some(
+        1,
+    ),
+}