Merge pull request #1903 from zed-industries/override-pyright-completion-sorting

Add LspAdapter hook for processing completions, fix completion sorting from Pyright

crates/language/src/language.rs

@@ -135,6 +135,10 @@ impl CachedLspAdapter {
         self.adapter.process_diagnostics(params).await
     }
 
+    pub async fn process_completion(&self, completion_item: &mut lsp::CompletionItem) {
+        self.adapter.process_completion(completion_item).await
+    }
+
     pub async fn label_for_completion(
         &self,
         completion_item: &lsp::CompletionItem,
@@ -175,6 +179,8 @@ pub trait LspAdapter: 'static + Send + Sync {
 
     async fn process_diagnostics(&self, _: &mut lsp::PublishDiagnosticsParams) {}
 
+    async fn process_completion(&self, _: &mut lsp::CompletionItem) {}
+
     async fn label_for_completion(
         &self,
         _: &lsp::CompletionItem,
@@ -826,6 +832,12 @@ impl Language {
         }
     }
 
+    pub async fn process_completion(self: &Arc<Self>, completion: &mut lsp::CompletionItem) {
+        if let Some(adapter) = self.adapter.as_ref() {
+            adapter.process_completion(completion).await;
+        }
+    }
+
     pub async fn label_for_completion(
         self: &Arc<Self>,
         completion: &lsp::CompletionItem,

crates/language/src/proto.rs

@@ -426,10 +426,11 @@ pub async fn deserialize_completion(
         .and_then(deserialize_anchor)
         .ok_or_else(|| anyhow!("invalid old end"))?;
     let lsp_completion = serde_json::from_slice(&completion.lsp_completion)?;
-    let label = match language {
-        Some(l) => l.label_for_completion(&lsp_completion).await,
-        None => None,
-    };
+
+    let mut label = None;
+    if let Some(language) = language {
+        label = language.label_for_completion(&lsp_completion).await;
+    }
 
     Ok(Completion {
         old_range: old_start..old_end,

crates/project/src/project.rs

@@ -3329,7 +3329,9 @@ impl Project {
                     let snapshot = this.snapshot();
                     let clipped_position = this.clip_point_utf16(position, Bias::Left);
                     let mut range_for_token = None;
-                    completions.into_iter().filter_map(move |lsp_completion| {
+                    completions
+                        .into_iter()
+                        .filter_map(move |mut lsp_completion| {
                             // For now, we can only handle additional edits if they are returned
                             // when resolving the completion, not if they are present initially.
                             if lsp_completion
@@ -3340,7 +3342,8 @@ impl Project {
                                 return None;
                             }
 
-                        let (old_range, mut new_text) = match lsp_completion.text_edit.as_ref() {
+                            let (old_range, mut new_text) = match lsp_completion.text_edit.as_ref()
+                            {
                                 // If the language server provides a range to overwrite, then
                                 // check that the range is valid.
                                 Some(lsp::CompletionTextEdit::Edit(edit)) => {
@@ -3393,11 +3396,11 @@ impl Project {
                             LineEnding::normalize(&mut new_text);
                             let language = language.clone();
                             Some(async move {
-                            let label = if let Some(language) = language {
-                                language.label_for_completion(&lsp_completion).await
-                            } else {
-                                None
-                            };
+                                let mut label = None;
+                                if let Some(language) = language {
+                                    language.process_completion(&mut lsp_completion).await;
+                                    label = language.label_for_completion(&lsp_completion).await;
+                                }
                                 Completion {
                                     old_range,
                                     new_text,

crates/zed/src/languages/python.rs

@@ -87,6 +87,25 @@ impl LspAdapter for PythonLspAdapter {
         .log_err()
     }
 
+    async fn process_completion(&self, item: &mut lsp::CompletionItem) {
+        // Pyright assigns each completion item a `sortText` of the form `XX.YYYY.name`.
+        // Where `XX` is the sorting category, `YYYY` is based on most recent usage,
+        // and `name` is the symbol name itself.
+        //
+        // Because the the symbol name is included, there generally are not ties when
+        // sorting by the `sortText`, so the symbol's fuzzy match score is not taken
+        // into account. Here, we remove the symbol name from the sortText in order
+        // to allow our own fuzzy score to be used to break ties.
+        //
+        // see https://github.com/microsoft/pyright/blob/95ef4e103b9b2f129c9320427e51b73ea7cf78bd/packages/pyright-internal/src/languageService/completionProvider.ts#LL2873
+        let Some(sort_text) = &mut item.sort_text else { return };
+        let mut parts = sort_text.split('.');
+        let Some(first) = parts.next() else { return };
+        let Some(second) = parts.next() else { return };
+        let Some(_) = parts.next() else { return };
+        sort_text.replace_range(first.len() + second.len() + 1.., "");
+    }
+
     async fn label_for_completion(
         &self,
         item: &lsp::CompletionItem,