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

Max Brunsfeld created

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

Change summary

crates/language/src/language.rs    |  12 ++
crates/language/src/proto.rs       |   9 +
crates/project/src/project.rs      | 155 ++++++++++++++++---------------
crates/zed/src/languages/python.rs |  19 +++
4 files changed, 115 insertions(+), 80 deletions(-)

Detailed changes

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,88 +3329,91 @@ 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| {
-                        // 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
-                            .additional_text_edits
-                            .as_ref()
-                            .map_or(false, |edits| !edits.is_empty())
-                        {
-                            return None;
-                        }
+                    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
+                                .additional_text_edits
+                                .as_ref()
+                                .map_or(false, |edits| !edits.is_empty())
+                            {
+                                return None;
+                            }
 
-                        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)) => {
-                                let range = range_from_lsp(edit.range);
-                                let start = snapshot.clip_point_utf16(range.start, Bias::Left);
-                                let end = snapshot.clip_point_utf16(range.end, Bias::Left);
-                                if start != range.start || end != range.end {
-                                    log::info!("completion out of expected range");
-                                    return None;
+                            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)) => {
+                                    let range = range_from_lsp(edit.range);
+                                    let start = snapshot.clip_point_utf16(range.start, Bias::Left);
+                                    let end = snapshot.clip_point_utf16(range.end, Bias::Left);
+                                    if start != range.start || end != range.end {
+                                        log::info!("completion out of expected range");
+                                        return None;
+                                    }
+                                    (
+                                        snapshot.anchor_before(start)..snapshot.anchor_after(end),
+                                        edit.new_text.clone(),
+                                    )
                                 }
-                                (
-                                    snapshot.anchor_before(start)..snapshot.anchor_after(end),
-                                    edit.new_text.clone(),
-                                )
-                            }
-                            // If the language server does not provide a range, then infer
-                            // the range based on the syntax tree.
-                            None => {
-                                if position != clipped_position {
-                                    log::info!("completion out of expected range");
+                                // If the language server does not provide a range, then infer
+                                // the range based on the syntax tree.
+                                None => {
+                                    if position != clipped_position {
+                                        log::info!("completion out of expected range");
+                                        return None;
+                                    }
+                                    let Range { start, end } = range_for_token
+                                        .get_or_insert_with(|| {
+                                            let offset = position.to_offset(&snapshot);
+                                            let (range, kind) = snapshot.surrounding_word(offset);
+                                            if kind == Some(CharKind::Word) {
+                                                range
+                                            } else {
+                                                offset..offset
+                                            }
+                                        })
+                                        .clone();
+                                    let text = lsp_completion
+                                        .insert_text
+                                        .as_ref()
+                                        .unwrap_or(&lsp_completion.label)
+                                        .clone();
+                                    (
+                                        snapshot.anchor_before(start)..snapshot.anchor_after(end),
+                                        text,
+                                    )
+                                }
+                                Some(lsp::CompletionTextEdit::InsertAndReplace(_)) => {
+                                    log::info!("unsupported insert/replace completion");
                                     return None;
                                 }
-                                let Range { start, end } = range_for_token
-                                    .get_or_insert_with(|| {
-                                        let offset = position.to_offset(&snapshot);
-                                        let (range, kind) = snapshot.surrounding_word(offset);
-                                        if kind == Some(CharKind::Word) {
-                                            range
-                                        } else {
-                                            offset..offset
-                                        }
-                                    })
-                                    .clone();
-                                let text = lsp_completion
-                                    .insert_text
-                                    .as_ref()
-                                    .unwrap_or(&lsp_completion.label)
-                                    .clone();
-                                (
-                                    snapshot.anchor_before(start)..snapshot.anchor_after(end),
-                                    text,
-                                )
-                            }
-                            Some(lsp::CompletionTextEdit::InsertAndReplace(_)) => {
-                                log::info!("unsupported insert/replace completion");
-                                return None;
-                            }
-                        };
-
-                        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
                             };
-                            Completion {
-                                old_range,
-                                new_text,
-                                label: label.unwrap_or_else(|| {
-                                    CodeLabel::plain(
-                                        lsp_completion.label.clone(),
-                                        lsp_completion.filter_text.as_deref(),
-                                    )
-                                }),
-                                lsp_completion,
-                            }
+
+                            LineEnding::normalize(&mut new_text);
+                            let language = language.clone();
+                            Some(async move {
+                                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,
+                                    label: label.unwrap_or_else(|| {
+                                        CodeLabel::plain(
+                                            lsp_completion.label.clone(),
+                                            lsp_completion.filter_text.as_deref(),
+                                        )
+                                    }),
+                                    lsp_completion,
+                                }
+                            })
                         })
-                    })
                 });
 
                 Ok(futures::future::join_all(completions).await)

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,