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/lspClient/languageClient.ts

@@ -1,14 +1,14 @@
+import { extname } from 'path'
 import { Uri, 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 { documentSelector, serverModule } from 'dvc-vscode-lsp'
 import { Disposable } from '../class/dispose'
-import { findFiles } from '../fileSystem/workspace'
 
 export class LanguageClientWrapper extends Disposable {
   private client: LanguageClient
@@ -35,33 +35,27 @@ 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
+    this.client.onRequest('isFile', (path: string) => {
+      const uri = Uri.file(path).toString()
+      const languageId = extname(path) === '.py' ? 'python' : 'other'
+      try {
+        const text = readFileSync(path, 'utf8')
+        return {
+          languageId,
+          text,
+          uri,
+          version: 0
+        }
+      } catch {
+        return null
       }
     })
 
-    await this.client.sendRequest('initialTextDocuments', {
-      textDocuments
-    })
+    void this.start()
+  }
 
-    return this
+  start() {
+    return this.client.start()
   }
 
   stop() {

languageServer/src/LanguageServer.ts

@@ -1,3 +1,4 @@
+import { dirname, join } from 'path'
 import {
   TextDocuments,
   InitializeResult,
@@ -11,15 +12,14 @@ import {
   Position,
   Range,
   DocumentSymbol,
-  TextDocumentItem
+  Connection
 } 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 +31,28 @@ 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
   ) {
     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)
@@ -132,7 +100,8 @@ export class LanguageServer {
     return locationsAccumulator
   }
 
-  private onDefinition(params: DefinitionParams) {
+  // eslint-disable-next-line sonarjs/cognitive-complexity
+  private async onDefinition(params: DefinitionParams, connection: Connection) {
     const document = this.getDvcTextDocument(params)
     const symbolUnderCursor = document?.symbolAt(params.position)
 
@@ -147,6 +116,33 @@ export class LanguageServer {
 
       locationsAccumulator.push(...fileLocations)
 
+      for (const possibleFile of symbolUnderCursor.name.split(' ')) {
+        const possiblePath = join(
+          dirname(document.uri).replace('file:///', ''),
+          possibleFile
+        )
+        const isFile = await connection.sendRequest<{
+          languageId: string
+          text: string
+          uri: URI
+          version: number
+        } | null>('isFile', possiblePath)
+        if (isFile) {
+          const { languageId, text, version } = isFile
+          const doc = TextDocument.create(
+            possiblePath,
+            languageId,
+            version,
+            text
+          )
+          const start = Position.create(0, 0)
+          const end = doc.positionAt(doc.getText().length - 1)
+          const range = Range.create(start, end)
+
+          locationsAccumulator.push(Location.create(possiblePath, range))
+        }
+      }
+
       const locationsFromOtherDocuments = this.getLocationsFromOtherDocuments(
         symbolUnderCursor,
         allDocs

languageServer/src/test/definitions.test.ts

@@ -1,17 +1,12 @@
 import { Position, Range } from 'vscode-languageserver/node'
-import {
-  foreach_dvc_yaml,
-  params_dvc_yaml,
-  vars_dvc_yaml
-} from './fixtures/examples/valid'
+import { foreach_dvc_yaml, params_dvc_yaml } from './fixtures/examples/valid'
 import { params } from './fixtures/params'
 import { requestDefinitions } from './utils/requestDefinitions'
 import { openTheseFilesAndNotifyServer } from './utils/openTheseFilesAndNotifyServer'
 import {
   disposeTestConnections,
   setupTestConnections
 } from './utils/setup-test-connections'
-import { sendTheseFilesToServer } from './utils/sendTheseFilesToServer'
 
 describe('textDocument/definitions', () => {
   beforeEach(() => {
@@ -74,26 +69,4 @@ describe('textDocument/definitions', () => {
       uri: 'file:///dvc.yaml'
     })
   })
-
-  it('should provide a single location that points to the top of the file path symbol', async () => {
-    const [dvcYaml] = await sendTheseFilesToServer([
-      {
-        languageId: 'yaml',
-        mockContents: vars_dvc_yaml,
-        mockPath: 'dvc.yaml'
-      },
-      {
-        languageId: 'json',
-        mockContents: '',
-        mockPath: 'params.json'
-      }
-    ])
-    const response = await requestDefinitions(dvcYaml, 'params.json')
-
-    expect(response).toBeTruthy()
-    expect(response).toStrictEqual({
-      range: Range.create(Position.create(0, 0), Position.create(0, 0)),
-      uri: 'file:///params.json'
-    })
-  })
 })

languageServer/src/test/fixtures/examples/valid/index.ts

@@ -151,32 +151,3 @@ stages:
       - output
 
 `
-export const vars_dvc_yaml = `
-vars:
-- custom_params.yaml
-- models:
-      us:
-        threshold: 10
-- desc: 'Reusable description'
-- params.json
-- myvar: 'value'
-- config/myapp.yaml
-- params.json:clean,feats
-
-stages:
-  test_vars:
-    vars:
-      - params.json:build
-      - model:
-          filename: 'model-us.hdf5'
-    cmd: echo hello world
-  test_foreach_vars:
-    foreach: \${vars}
-    do:
-      vars:
-        - params.json:build
-        - model:
-            filename: 'model-us.hdf5'
-      cmd: echo \${item} \${model.filename}
-
-`

languageServer/src/test/utils/setup-test-connections.ts

@@ -22,6 +22,7 @@ export const setupTestConnections = () => {
 
   server = createConnection(up, down)
   client = createConnection(down, up)
+  client.onRequest('isFile', () => null)
 
   dvcLanguageService.listen(server)
   client.listen()