Implement simple onDefinition feature in LSP

extension/src/extension.ts

@@ -50,6 +50,7 @@ import { Setup } from './setup'
 import { definedAndNonEmpty } from './util/array'
 import { stopProcesses } from './processExecution'
 import { Flag } from './cli/dvc/constants'
+import { LanguageClientWrapper } from './lspClient/languageClient'
 
 export class Extension extends Disposable {
   protected readonly internalCommands: InternalCommands
@@ -255,6 +256,8 @@ export class Extension extends Disposable {
 
     void showWalkthroughOnFirstUse(env.isNewAppInstall)
     this.dispose.track(recommendRedHatExtensionOnce())
+
+    this.dispose.track(new LanguageClientWrapper())
   }
 
   public async initialize() {

extension/src/fileSystem/index.ts

@@ -204,3 +204,15 @@ export const getBinDisplayText = (
     ? '.' + sep + relative(workspaceRoot, path)
     : path
 }
+
+export const readFileContents = (
+  uriString: string
+): { contents: string } | null => {
+  try {
+    const uri = Uri.parse(uriString)
+    if (!isDirectory(uri.fsPath)) {
+      return { contents: readFileSync(uri.fsPath, 'utf8') }
+    }
+  } catch {}
+  return null
+}

extension/src/lspClient/languageClient.ts

@@ -1,14 +1,13 @@
-import { Uri, workspace } from 'vscode'
+import { workspace } from 'vscode'
 import {
   LanguageClient,
   LanguageClientOptions,
   ServerOptions,
   TransportKind
 } from 'vscode-languageclient/node'
 import { documentSelector, serverModule } from 'dvc-vscode-lsp'
-import { readFileSync } from 'fs-extra'
 import { Disposable } from '../class/dispose'
-import { findFiles } from '../fileSystem/workspace'
+import { readFileContents } from '../fileSystem'
 
 export class LanguageClientWrapper extends Disposable {
   private client: LanguageClient
@@ -20,9 +19,7 @@ export class LanguageClientWrapper extends Disposable {
       documentSelector,
 
       synchronize: {
-        fileEvents: workspace.createFileSystemWatcher(
-          '**/*.{yaml,dvc,dvc.lock,json,toml}'
-        )
+        fileEvents: workspace.createFileSystemWatcher('**/dvc.yaml')
       }
     }
 
@@ -35,37 +32,11 @@ export class LanguageClientWrapper extends Disposable {
       )
     )
 
-    // Start the client. This will also launch the server
-    void this.start()
-  }
-
-  async start() {
-    await this.client.start()
-
-    const files = await findFiles('**/*.{yaml,json,py,toml}', '.??*')
-
-    const textDocuments = files.map(filePath => {
-      const uri = Uri.file(filePath).toString()
-      const languageId = filePath.endsWith('yaml') ? 'yaml' : 'json'
-      const text = readFileSync(filePath, 'utf8')
-
-      return {
-        languageId,
-        text,
-        uri,
-        version: 0
-      }
-    })
-
-    await this.client.sendRequest('initialTextDocuments', {
-      textDocuments
-    })
-
-    return this
-  }
+    this.dispose.track(
+      this.client.onRequest('readFileContents', readFileContents)
+    )
 
-  stop() {
-    void this.client.stop()
+    void this.client.start()
   }
 
   private getServerOptions(): ServerOptions {

extension/src/test/suite/fileSystem/index.test.ts

@@ -1,12 +1,18 @@
 import { join, resolve } from 'path'
 import process from 'process'
+import { Uri } from 'vscode'
 import { afterEach, beforeEach, describe, it, suite } from 'mocha'
 import { expect } from 'chai'
 import { Disposable } from '@hediet/std/disposable'
 import { restore } from 'sinon'
 import { ensureFileSync, removeSync, writeFileSync } from 'fs-extra'
 import { dvcDemoPath } from '../../util'
-import { checkSignalFile, exists, getGitPath } from '../../../fileSystem'
+import {
+  checkSignalFile,
+  exists,
+  getGitPath,
+  readFileContents
+} from '../../../fileSystem'
 import { gitPath } from '../../../cli/git/constants'
 import { GitReader } from '../../../cli/git/reader'
 import { standardizePath } from '../../../fileSystem/path'
@@ -93,4 +99,20 @@ suite('File System Test Suite', () => {
       expect(exists(mockSignalFilePath)).to.be.false
     })
   })
+
+  describe('readFileContents', () => {
+    it('should read the contents of a file when it exists', () => {
+      const uriString = Uri.file(join(dvcDemoPath, 'train.py')).toString()
+
+      const file = readFileContents(uriString)
+      expect(file?.contents).to.contain('main()')
+    })
+
+    it('should return null when the file cannot be found', () => {
+      const uriString = 'file:///some/fun/file.txt'
+
+      const contents = readFileContents(uriString)
+      expect(contents).to.be.null
+    })
+  })
 })

languageServer/src/LanguageServer.ts

@@ -1,3 +1,4 @@
+import { dirname } from 'path'
 import {
   TextDocuments,
   InitializeResult,
@@ -7,19 +8,16 @@ import {
   CodeActionParams,
   DefinitionParams,
   SymbolKind,
-  Location,
-  Position,
-  Range,
   DocumentSymbol,
-  TextDocumentItem
+  Connection,
+  Location
 } from 'vscode-languageserver/node'
 import { TextDocument } from 'vscode-languageserver-textdocument'
 import { URI } from 'vscode-uri'
 import { TextDocumentWrapper } from './TextDocumentWrapper'
 
 export class LanguageServer {
   private documentsKnownToEditor!: TextDocuments<TextDocument>
-  private documentsFromDvcClient: TextDocumentWrapper[] = []
 
   public listen(connection: _Connection) {
     this.documentsKnownToEditor = new TextDocuments(TextDocument)
@@ -31,60 +29,26 @@ export class LanguageServer {
         return null
       }
 
-      return this.onDefinition(params)
+      return this.onDefinition(params, connection)
     })
 
-    connection.onRequest(
-      'initialTextDocuments',
-      (params: { textDocuments: TextDocumentItem[] }) => {
-        this.documentsFromDvcClient = params.textDocuments.map(
-          ({ uri, languageId, version, text: content }) => {
-            const textDocument = TextDocument.create(
-              uri,
-              languageId,
-              version,
-              content
-            )
-
-            return this.wrap(textDocument)
-          }
-        )
-      }
-    )
-
     this.documentsKnownToEditor.listen(connection)
 
     connection.listen()
   }
 
   private getAllDocuments() {
     const openDocuments = this.documentsKnownToEditor.all()
-    const acc: TextDocumentWrapper[] = openDocuments.map(doc => this.wrap(doc))
-
-    for (const textDocument of this.documentsFromDvcClient) {
-      const userAlreadyOpenedIt = this.documentsKnownToEditor.get(
-        textDocument.uri
-      )
-
-      if (!userAlreadyOpenedIt) {
-        acc.push(textDocument)
-      }
-    }
-
-    return acc
+    return openDocuments.map(doc => this.wrap(doc))
   }
 
-  private getDvcTextDocument(
-    params: TextDocumentPositionParams | CodeActionParams
-  ) {
+  private getDocument(params: TextDocumentPositionParams | CodeActionParams) {
     const uri = params.textDocument.uri
+
     const doc = this.documentsKnownToEditor.get(uri)
 
     if (!doc) {
-      const alternative = this.documentsFromDvcClient.find(
-        txtDoc => txtDoc.uri === uri
-      )
-      return alternative ?? null
+      return null
     }
 
     return this.wrap(doc)
@@ -94,78 +58,70 @@ export class LanguageServer {
     return new TextDocumentWrapper(doc)
   }
 
-  private getFilePathLocations(
-    symbolUnderCursor: DocumentSymbol,
-    allDocs: TextDocumentWrapper[]
-  ) {
+  private getKnownDocumentLocations(symbolUnderCursor: DocumentSymbol) {
     if (symbolUnderCursor.kind !== SymbolKind.File) {
       return []
     }
 
     const filePath = symbolUnderCursor.name
 
-    const matchingFiles = allDocs.filter(doc =>
-      URI.file(doc.uri).fsPath.endsWith(filePath)
+    const matchingFiles = this.getAllDocuments().filter(doc =>
+      URI.parse(doc.uri).fsPath.endsWith(filePath)
     )
 
-    return matchingFiles.map(doc => {
-      const uri = doc.uri
-      const start = Position.create(0, 0)
-      const end = doc.positionAt(doc.getText().length - 1)
-      const range = Range.create(start, end)
-
-      return Location.create(uri, range)
-    })
+    return matchingFiles.map(doc => doc.getLocation())
   }
 
-  private getLocationsFromOtherDocuments(
-    symbolUnderCursor: DocumentSymbol,
-    allDocs: TextDocumentWrapper[]
-  ) {
-    const locationsAccumulator = []
-
-    for (const txtDoc of allDocs) {
-      const locations = txtDoc.findLocationsFor(symbolUnderCursor)
-      locationsAccumulator.push(...locations)
-    }
+  private async onDefinition(params: DefinitionParams, connection: Connection) {
+    const document = this.getDocument(params)
 
-    return locationsAccumulator
-  }
-
-  private onDefinition(params: DefinitionParams) {
-    const document = this.getDvcTextDocument(params)
     const symbolUnderCursor = document?.symbolAt(params.position)
 
-    if (document && symbolUnderCursor) {
-      const allDocs = this.getAllDocuments()
-      const locationsAccumulator = []
-
-      const fileLocations = this.getFilePathLocations(
-        symbolUnderCursor,
-        allDocs
-      )
+    if (!(document && symbolUnderCursor)) {
+      return null
+    }
 
-      locationsAccumulator.push(...fileLocations)
+    const fileLocations = this.getKnownDocumentLocations(symbolUnderCursor)
 
-      const locationsFromOtherDocuments = this.getLocationsFromOtherDocuments(
+    if (fileLocations.length === 0) {
+      await this.checkIfSymbolsAreFiles(
+        connection,
+        document,
         symbolUnderCursor,
-        allDocs
+        fileLocations
       )
+    }
 
-      locationsAccumulator.push(...locationsFromOtherDocuments)
+    if (fileLocations.length > 0) {
+      return this.arrayOrSingleResponse(fileLocations)
+    }
 
-      const externalLocations = locationsAccumulator.filter(
-        location => location.uri !== document.uri
-      )
+    return null
+  }
 
-      if (externalLocations.length > 0) {
-        return this.arrayOrSingleResponse(externalLocations)
+  private async checkIfSymbolsAreFiles(
+    connection: _Connection,
+    document: TextDocumentWrapper,
+    symbolUnderCursor: DocumentSymbol,
+    fileLocations: Location[]
+  ) {
+    for (const possibleFile of symbolUnderCursor.name.split(' ')) {
+      const possiblePath = URI.parse(
+        [dirname(document.uri), possibleFile].join('/')
+      ).toString()
+      const file = await connection.sendRequest<{
+        contents: string
+      } | null>('readFileContents', possiblePath)
+      if (file) {
+        const location = this.getLocation(possiblePath, file.contents)
+        fileLocations.push(location)
       }
-
-      return this.arrayOrSingleResponse(locationsAccumulator)
     }
+  }
 
-    return null
+  private getLocation(path: string, contents: string) {
+    const doc = this.wrap(TextDocument.create(path, 'plain/text', 0, contents))
+    return doc.getLocation()
   }
 
   private arrayOrSingleResponse<T>(elements: T[]) {