Fix/improve completion sorting

krassowski · krassowski · commit 6eacd1048569 · 2023-09-16T16:30:18.000+01:00
diff --git a/packages/jupyterlab-lsp/src/features/completion/item.ts b/packages/jupyterlab-lsp/src/features/completion/item.ts
@@ -30,10 +30,6 @@ namespace CompletionItem {
 }
 
 export class CompletionItem implements IExtendedCompletionItem {
-  private _detail: string | undefined;
-  private _documentation: string | undefined;
-  private _isDocumentationMarkdown: boolean;
-  private _resolved: boolean;
   /**
    * Self-reference to make sure that the instance for will remain accessible
    * after any copy operation (whether via spread syntax or Object.assign)
@@ -42,7 +38,6 @@ export class CompletionItem implements IExtendedCompletionItem {
   public self: CompletionItem;
   public element: HTMLLIElement;
   public source: string;
-  private _currentInsertText: string;
 
   get isDocumentationMarkdown(): boolean {
     return this._isDocumentationMarkdown;
@@ -55,15 +50,14 @@ export class CompletionItem implements IExtendedCompletionItem {
   public label: string;
 
   icon: LabIcon | undefined;
-  private match: lsProtocol.CompletionItem;
 
   constructor(protected options: CompletionItem.IOptions) {
     const match = options.match;
     this.label = match.label;
     this._setDocumentation(match.documentation);
     this._resolved = false;
     this._detail = match.detail;
-    this.match = match;
+    this._match = match;
     this.self = this;
     this.source = options.source;
     this.icon = options.icon ? options.icon : undefined;
@@ -74,48 +68,38 @@ export class CompletionItem implements IExtendedCompletionItem {
     // Ideally this would be fixed and tested e2e in JupyterLab 4.0.7.
     // https://github.com/jupyterlab/jupyterlab/issues/15125
     makeGetterEnumerable(this, 'insertText');
+    makeGetterEnumerable(this, 'sortText');
+    makeGetterEnumerable(this, 'filterText');
   }
 
   get type() {
     return this.options.type;
   }
 
-  private _setDocumentation(
-    documentation: string | lsProtocol.MarkupContent | undefined
-  ) {
-    if (lsProtocol.MarkupContent.is(documentation)) {
-      this._documentation = documentation.value;
-      this._isDocumentationMarkdown = documentation.kind === 'markdown';
-    } else {
-      this._documentation = documentation;
-      this._isDocumentationMarkdown = false;
-    }
-  }
-
   /**
    * Completion to be inserted.
    */
   get insertText(): string {
-    return this._currentInsertText || this.match.insertText || this.match.label;
+    return (
+      this._currentInsertText || this._match.insertText || this._match.label
+    );
   }
-
   set insertText(text: string) {
     this._currentInsertText = text;
   }
 
   get sortText(): string {
-    return this.match.sortText || this.match.label;
+    return this._currentSortText || this._match.sortText || this._match.label;
+  }
+  set sortText(text: string) {
+    this._currentSortText = text;
   }
 
   get filterText(): string | undefined {
-    return this.match.filterText;
+    return this._match.filterText;
   }
-
-  private _supportsResolution(): boolean {
-    const connection = this.options.connection;
-    return (
-      connection.serverCapabilities.completionProvider?.resolveProvider ?? false
-    );
+  set filterText(text: string | undefined) {
+    this._match.filterText = text;
   }
 
   get detail(): string | undefined {
@@ -153,7 +137,7 @@ export class CompletionItem implements IExtendedCompletionItem {
 
     const resolvedCompletionItem = await connection.clientRequests[
       'completionItem/resolve'
-    ].request(this.match);
+    ].request(this._match);
 
     if (resolvedCompletionItem === null) {
       return this;
@@ -181,16 +165,43 @@ export class CompletionItem implements IExtendedCompletionItem {
    * Indicates if the item is deprecated.
    */
   get deprecated(): boolean {
-    if (this.match.deprecated) {
-      return this.match.deprecated;
+    if (this._match.deprecated) {
+      return this._match.deprecated;
     }
     return (
-      this.match.tags != null &&
-      this.match.tags.some(
+      this._match.tags != null &&
+      this._match.tags.some(
         tag => tag == lsProtocol.CompletionItemTag.Deprecated
       )
     );
   }
+
+  private _setDocumentation(
+    documentation: string | lsProtocol.MarkupContent | undefined
+  ) {
+    if (lsProtocol.MarkupContent.is(documentation)) {
+      this._documentation = documentation.value;
+      this._isDocumentationMarkdown = documentation.kind === 'markdown';
+    } else {
+      this._documentation = documentation;
+      this._isDocumentationMarkdown = false;
+    }
+  }
+
+  private _supportsResolution(): boolean {
+    const connection = this.options.connection;
+    return (
+      connection.serverCapabilities.completionProvider?.resolveProvider ?? false
+    );
+  }
+
+  private _detail: string | undefined;
+  private _documentation: string | undefined;
+  private _isDocumentationMarkdown: boolean;
+  private _resolved: boolean;
+  private _currentInsertText: string;
+  private _currentSortText: string;
+  private _match: lsProtocol.CompletionItem;
 }
 
 function makeGetterEnumerable(instance: object, name: string) {
diff --git a/packages/jupyterlab-lsp/src/features/completion/model.spec.ts b/packages/jupyterlab-lsp/src/features/completion/model.spec.ts
@@ -8,14 +8,15 @@ describe('LSPCompleterModel', () => {
 
   function createDummyItem(
     match: lsProtocol.CompletionItem,
-    type: string = 'dummy'
+    type: string = 'dummy',
+    source: string = 'lsp'
   ) {
     return new CompletionItem({
       type,
       icon: null as any,
       match,
       connection: null as any,
-      source: 'lsp'
+      source: source
     });
   }
 
@@ -79,6 +80,68 @@ describe('LSPCompleterModel', () => {
     expect(sortedItems.map(item => item.sortText)).toEqual(['a', 'b', 'c']);
   });
 
+  describe('sorting by source', () => {
+    const testCompletionA = createDummyItem(
+      {
+        label: 'test'
+      },
+      'a',
+      'LSP'
+    );
+    const testCompletionB = createDummyItem(
+      {
+        label: 'test'
+      },
+      'b',
+      'kernel'
+    );
+    const testCompletionC = createDummyItem(
+      {
+        label: 'test'
+      },
+      'c',
+      'context'
+    );
+    const testCompletionD = createDummyItem(
+      {
+        label: 'test'
+      },
+      'd',
+      'unknown'
+    );
+    const completionsFromDifferentSources = [
+      testCompletionA,
+      testCompletionC,
+      testCompletionB
+    ];
+
+    it('completions are sorted by source', () => {
+      model.setCompletionItems(completionsFromDifferentSources);
+      model.query = 'test';
+      let sortedItems = model.completionItems();
+      expect(sortedItems.map(item => item.type)).toEqual(['a', 'b', 'c']);
+    });
+
+    it('kernel completions can be prioritised', () => {
+      model.setCompletionItems(completionsFromDifferentSources);
+      model.query = 'test';
+      model.settings.kernelCompletionsFirst = true;
+      let sortedItems = model.completionItems();
+      expect(sortedItems.map(item => item.type)).toEqual(['b', 'a', 'c']);
+    });
+
+    it('completions from unknown source land at the end', () => {
+      model.setCompletionItems([
+        testCompletionD,
+        ...completionsFromDifferentSources
+      ]);
+      model.query = 'test';
+      let sortedItems = model.completionItems();
+      const types = sortedItems.map(item => item.type);
+      expect(types[types.length - 1]).toEqual('d');
+    });
+  });
+
   it('ignores perfect matches when asked', () => {
     model = new LSPCompleterModel({
       includePerfectMatches: false
diff --git a/packages/jupyterlab-lsp/src/features/completion/model.ts b/packages/jupyterlab-lsp/src/features/completion/model.ts
@@ -38,14 +38,16 @@ export class GenericCompleterModel<
 
   completionItems(): T[] {
     let query = this.query;
-    // TODO: make use of `processedItemsCache` when made available upstream,
-    // see https://github.com/jupyterlab/jupyterlab/pull/15025
     // (setting query is bad because it resets the cache; ideally we would
     // modify the sorting and filtering algorithm upstream).
+
+    // TODO processedItemsCache
     this.query = '';
+
     let unfilteredItems = (
       super.completionItems() as CompletionHandler.ICompletionItem[]
     ).map(this.harmoniseItem);
+
     this.query = query;
 
     // always want to sort
@@ -178,7 +180,7 @@ export class GenericCompleterModel<
       }
     }
 
-    results.sort(this.compareMatches);
+    results.sort(this.compareMatches.bind(this));
 
     return results.map(x => x.item);
   }
@@ -206,14 +208,19 @@ export namespace GenericCompleterModel {
      */
     includePerfectMatches?: boolean;
     /**
-     * Wheteher matches should be pre-filtered (default = true)
+     * Whether matches should be pre-filtered (default = true)
      */
     preFilterMatches?: boolean;
+    /**
+     * Whether kernel completions should be shown first.
+     */
+    kernelCompletionsFirst?: boolean;
   }
   export const defaultOptions: IOptions = {
     caseSensitive: true,
     includePerfectMatches: true,
-    preFilterMatches: true
+    preFilterMatches: true,
+    kernelCompletionsFirst: false
   };
 }
 
@@ -239,13 +246,22 @@ export class LSPCompleterModel extends GenericCompleterModel<MaybeCompletionItem
     a: ICompletionMatch<MaybeCompletionItem>,
     b: ICompletionMatch<MaybeCompletionItem>
   ): number {
-    const delta = a.score - b.score;
-    if (delta !== 0) {
-      return delta;
-    }
-    // solve ties using sortText
-
-    // note: locale compare is case-insensitive
-    return (a.item.sortText ?? 'z').localeCompare(b.item.sortText ?? 'z');
+    // TODO: take source order from provider ranks, upstream this code
+    const sourceOrder = {
+      LSP: 1,
+      kernel: this.settings.kernelCompletionsFirst ? 0 : 2,
+      context: 3
+    };
+    const aRank = a.item.source
+      ? sourceOrder[a.item.source as keyof typeof sourceOrder] ?? 4
+      : 4;
+    const bRank = b.item.source
+      ? sourceOrder[b.item.source as keyof typeof sourceOrder] ?? 4
+      : 4;
+    return (
+      aRank - bRank ||
+      (a.item.sortText ?? 'z').localeCompare(b.item.sortText ?? 'z') ||
+      a.score - b.score
+    );
   }
 }
diff --git a/packages/jupyterlab-lsp/src/features/completion/overrides.ts b/packages/jupyterlab-lsp/src/features/completion/overrides.ts
@@ -36,21 +36,17 @@ export class EnhancedContextCompleterProvider extends ContextCompleterProvider {
     super();
   }
 
-  private get _kernelCompletionsFirst(): boolean {
-    return this.options.settings.composite.kernelCompletionsFirst;
-  }
-
   async fetch(
     request: CompletionHandler.IRequest,
     context: ICompletionContext
   ): Promise<CompletionHandler.ICompletionItemsReply> {
     const result = await super.fetch(request, context);
-    result.items = result.items.map(i => {
+    result.items = result.items.map((item, order) => {
       return {
-        ...i,
-        icon: this.iconFor(i.type ?? 'Text') ?? this.iconFor('Text'),
-        type: i.type === '<unknown>' ? undefined : (i.type as string),
-        sortText: this._kernelCompletionsFirst ? 'a' : 'z',
+        ...item,
+        icon: this.iconFor(item.type ?? 'Text') ?? this.iconFor('Text'),
+        type: item.type === '<unknown>' ? undefined : (item.type as string),
+        sortText: String.fromCharCode(order),
         source: this.label
       };
     });
@@ -75,11 +71,11 @@ export class EnhancedKernelCompleterProvider extends KernelCompleterProvider {
     context: ICompletionContext
   ): Promise<CompletionHandler.ICompletionItemsReply> {
     const result = await super.fetch(request, context);
-    result.items = result.items.map(i => {
+    result.items = result.items.map((item, order) => {
       return {
-        ...i,
-        icon: this.iconFor(i.type ?? KernelKind) ?? this.iconFor(KernelKind),
-        sortText: 'z',
+        ...item,
+        icon: this.iconFor(item.type ?? KernelKind) ?? this.iconFor(KernelKind),
+        sortText: String.fromCharCode(order),
         source: this.label
       };
     });
@@ -88,7 +84,7 @@ export class EnhancedKernelCompleterProvider extends KernelCompleterProvider {
 
   async isApplicable(context: ICompletionContext): Promise<boolean> {
     // Note: this method logs errors instead of throwing to ensure we do not ever
-    // break the upstream kernel completer, even if there is an error elsehwere.
+    // break the upstream kernel completer, even if there is an error elsewhere.
     const upstream = await super.isApplicable(context);
 
     if (upstream === false) {
diff --git a/packages/jupyterlab-lsp/src/features/completion/provider.ts b/packages/jupyterlab-lsp/src/features/completion/provider.ts
@@ -65,13 +65,15 @@ export class CompletionProvider implements ICompletionProvider<CompletionItem> {
     const model = new LSPCompleterModel({
       caseSensitive: composite.caseSensitive,
       preFilterMatches: composite.preFilterMatches,
-      includePerfectMatches: composite.includePerfectMatches
+      includePerfectMatches: composite.includePerfectMatches,
+      kernelCompletionsFirst: composite.kernelCompletionsFirst
     });
     this.options.settings.changed.connect(() => {
       const composite = this.options.settings.composite;
       model.settings.caseSensitive = composite.caseSensitive;
       model.settings.preFilterMatches = composite.preFilterMatches;
       model.settings.includePerfectMatches = composite.includePerfectMatches;
+      model.settings.kernelCompletionsFirst = composite.kernelCompletionsFirst;
     });
     return model;
   };
