save work

languageServer/src/TextDocumentWrapper.ts

@@ -16,6 +16,8 @@ import {
   Pair,
   isPair
 } from 'yaml'
+import { has } from 'lodash'
+import { findNodeAtLocation, parse, parseTree } from 'jsonc-parser'
 import * as RegExes from './regexes'
 import { ITextDocumentWrapper } from './ITextDocumentWrapper'
 
@@ -45,7 +47,16 @@ export class TextDocumentWrapper implements ITextDocumentWrapper {
     return parseDocument(this.getText())
   }
 
-  public findLocationsFor(symbol: DocumentSymbol) {
+  public findLocationsFor(symbol: DocumentSymbol): Location[] {
+    const propertyPath = symbol.kind === SymbolKind.Property && symbol.detail
+    const itIsHere = propertyPath && this.hasProperty(propertyPath)
+
+    if (itIsHere) {
+      const location = this.getPropertyLocation(propertyPath)
+
+      return location ? [location] : []
+    }
+
     const parts = symbol.name.split(/\s/g)
     const txt = this.getText()
 
@@ -69,6 +80,53 @@ export class TextDocumentWrapper implements ITextDocumentWrapper {
     return this.symbolScopeAt(position).pop()
   }
 
+  private getSymbolsFromPropertyPath(pathSegment: string, startIndex: number) {
+    const templateSymbols: DocumentSymbol[] = []
+    const symbols = pathSegment.matchAll(RegExes.alphadecimalWords)
+
+    const jsonPath: string[] = [] // Safe to assume, based on https://dvc.org/doc/user-guide/project-structure/dvcyaml-files#vars
+
+    for (const templateSymbol of symbols) {
+      const symbolName = templateSymbol[0]
+      const symbolJsonPath = [...jsonPath, symbolName]
+      const symbolStart = (templateSymbol.index ?? 0) + startIndex
+      const symbolEnd = symbolStart + templateSymbol[0].length
+      const symbolRange = Range.create(
+        this.positionAt(symbolStart),
+        this.positionAt(symbolEnd)
+      )
+
+      templateSymbols.push(
+        DocumentSymbol.create(
+          templateSymbol[0],
+          symbolJsonPath.join('.'),
+          SymbolKind.Property,
+          symbolRange,
+          symbolRange
+        )
+      )
+
+      jsonPath.push(symbolName)
+    }
+
+    return templateSymbols
+  }
+
+  private extractPropertyPathSymbolsFrom(text: string, startIndex: number) {
+    const symbols: DocumentSymbol[] = []
+    const pathLikeSegments = text.matchAll(RegExes.propertyPathLike)
+
+    for (const path of pathLikeSegments) {
+      const matchIndex = path.index ?? 0
+
+      symbols.push(
+        ...this.getSymbolsFromPropertyPath(path[0], startIndex + matchIndex)
+      )
+    }
+
+    return symbols
+  }
+
   private getTemplateExpressionSymbolsInsideScalar(
     scalarValue: string,
     nodeOffset: number
@@ -81,32 +139,10 @@ export class TextDocumentWrapper implements ITextDocumentWrapper {
       const expression = match[1]
       const matchOffset = match.index || 0
       const expressionOffset: number = nodeOffset + matchOffset + 2 // To account for the '${'
-      const symbols = expression.matchAll(RegExes.alphadecimalWords)
-
-      const jsonPath: string[] = [] // Safe to assume, based on https://dvc.org/doc/user-guide/project-structure/dvcyaml-files#vars
-
-      for (const templateSymbol of symbols) {
-        const symbolName = templateSymbol[0]
-        const symbolJsonPath = [...jsonPath, symbolName]
-        const symbolStart = (templateSymbol.index ?? 0) + expressionOffset
-        const symbolEnd = symbolStart + templateSymbol[0].length
-        const symbolRange = Range.create(
-          this.positionAt(symbolStart),
-          this.positionAt(symbolEnd)
-        )
-
-        templateSymbols.push(
-          DocumentSymbol.create(
-            templateSymbol[0],
-            symbolJsonPath.join('.'),
-            SymbolKind.Property,
-            symbolRange,
-            symbolRange
-          )
-        )
 
-        jsonPath.push(symbolName)
-      }
+      templateSymbols.push(
+        ...this.extractPropertyPathSymbolsFrom(expression, expressionOffset)
+      )
     }
 
     return templateSymbols
@@ -124,6 +160,22 @@ export class TextDocumentWrapper implements ITextDocumentWrapper {
       symbolKind = SymbolKind.File
     }
 
+    const children: DocumentSymbol[] = []
+
+    const variableTemplateSymbols = [
+      ...this.getTemplateExpressionSymbolsInsideScalar(nodeValue, nodeStart)
+    ]
+
+    if (variableTemplateSymbols.length > 0) {
+      children.push(...variableTemplateSymbols)
+    } else {
+      const propertyPathSymbols = this.extractPropertyPathSymbolsFrom(
+        nodeValue,
+        nodeStart
+      )
+      children.push(...propertyPathSymbols)
+    }
+
     const symbolsSoFar: DocumentSymbol[] = [
       DocumentSymbol.create(
         nodeValue,
@@ -132,7 +184,7 @@ export class TextDocumentWrapper implements ITextDocumentWrapper {
         Range.create(this.positionAt(nodeStart), this.positionAt(nodeEnd)),
         Range.create(this.positionAt(nodeStart), this.positionAt(valueEnd))
       ),
-      ...this.getTemplateExpressionSymbolsInsideScalar(nodeValue, nodeStart)
+      ...children
     ]
 
     return symbolsSoFar
@@ -158,19 +210,22 @@ export class TextDocumentWrapper implements ITextDocumentWrapper {
 
     const symbolsFound: Array<DocumentSymbol | null> = []
 
-    visit(this.getYamlDocument(), (_, node) => {
-      if (isNode(node) && node.range) {
-        const range = node.range
-        const nodeStart = range[0]
-        const nodeEnd = range[2]
-        const isCursorInsideNode =
-          cursorOffset >= nodeStart && cursorOffset <= nodeEnd
-
-        if (isCursorInsideNode) {
-          symbolsFound.push(...this.yamlNodeToDocumentSymbols(node, range))
+    if (this.uri.endsWith('yaml')) {
+      visit(this.getYamlDocument(), (_, node) => {
+        if (isNode(node) && node.range) {
+          const range = node.range
+          const nodeStart = range[0]
+          const nodeEnd = range[2]
+          const isCursorInsideNode =
+            cursorOffset >= nodeStart && cursorOffset <= nodeEnd
+
+          if (isCursorInsideNode) {
+            symbolsFound.push(...this.yamlNodeToDocumentSymbols(node, range))
+          }
         }
-      }
-    })
+      })
+    }
+
     const symbolStack = (symbolsFound.filter(Boolean) as DocumentSymbol[]).sort(
       (a, b) => {
         const offA = this.offsetAt(a.range.end) - this.offsetAt(a.range.start)
@@ -182,4 +237,59 @@ export class TextDocumentWrapper implements ITextDocumentWrapper {
 
     return [...symbolStack]
   }
+
+  private hasProperty(path: string) {
+    const parsedObj = this.toJSON()
+
+    return has(parsedObj, path)
+  }
+
+  private getPropertyLocation(path: string) {
+    const pathArray = path.split('.')
+
+    if (this.uri.endsWith('yaml')) {
+      const node = this.getYamlDocument().getIn(pathArray, true)
+
+      if (isNode(node) && node.range) {
+        const [nodeStart, , nodeEnd] = node.range
+        const start = this.positionAt(nodeStart)
+        const end = this.positionAt(nodeEnd)
+        const range = Range.create(start, end)
+        return Location.create(this.uri, range)
+      }
+
+      return null
+    }
+
+    if (this.uri.endsWith('json')) {
+      const rootNode = parseTree(this.getText())
+      const node = rootNode && findNodeAtLocation(rootNode, pathArray)
+
+      if (!node) {
+        return null
+      }
+      const nodeSrcIndex = node.offset
+      const nodeSrcLength = node.length
+      const nodeEnd = nodeSrcIndex + nodeSrcLength
+      const start = this.positionAt(nodeSrcIndex)
+      const end = this.positionAt(nodeEnd)
+      const range = Range.create(start, end)
+      return Location.create(this.uri, range)
+    }
+
+    return null
+  }
+
+  private toJSON() {
+    if (this.uri.endsWith('yaml')) {
+      return this.getYamlDocument().toJS()
+    }
+
+    if (this.uri.endsWith('json')) {
+      const src = this.getText()
+      return parse(src)
+    }
+
+    return null
+  }
 }

languageServer/src/regexes.ts

@@ -1,3 +1,4 @@
 export const variableTemplates = /\${([^}]+)}/g
 export const filePaths = /[\d/A-Za-z]+\.[A-Za-z]+/g
 export const alphadecimalWords = /[\dA-Za-z]+/g
+export const propertyPathLike = /[\d.A-Z[\]a-z]+/g