TargetElementEvaluator and wrapping PSI elements

IntelliJ Platform
1

Hi there forum :wave:
I need to turn YAML scalars into references. Consider this example:

automation:
  - alias: automation_name_1
    trigger:
       [...]

I wanted “Find usages” and the like from “automation_name_1” to work, but IntelliJ was complaining that I “Cannot search for usages from this location. Place the caret on the element to find usages for and try again.”, despite references from other files is working (if I use the notation automation.automation_name_1 in any other YAML file, Ctrl+Click correctly brings me to the example I wrote above).

I’m guessing this is happening because YAMLScalar is not a PsiNamedElement (right?), so I looked for a way of turning it into a named element. I found two possible ways:

  1. override the parser definition, but it was deemed not recommended - and it caused constant 100% CPU usage for some reason
  2. use TargetElementEvaluator and implement getNamedElement and adjustReferenceOrReferencedElement to return a wrapped element (implementation here)

I’m currently using approach 2, which leads me to the wrapper class itself:

/**
 * Wraps the "alias" key value element of an automation, giving the element [PsiNamedElement] features.
 */
class HassAutomation(private val element: YAMLScalar) : PsiNamedElement, YAMLScalar by element {

    /**
     * Automation name is actually the value of the "alias" key.
     */
    override fun getName(): String {
        return element.textValue
    }

    /**
     * Copied from [YAMLKeyValueImpl] + [org.jetbrains.yaml.YAMLUtil], although it seems like overkill.
     */
    override fun setName(name: String): PsiElement {
        if (name == element.textValue) {
            throw IncorrectOperationException(YAMLBundle.message("rename.same.name"))
        }

        val elementGenerator = YAMLElementGenerator.getInstance(element.project)

        val tempFile: PsiFile = elementGenerator.createDummyYamlWithText(name)
        val textElement = PsiTreeUtil.collectElementsOfType(tempFile, YAMLScalar::class.java).iterator().next()

        element.replace(textElement)
        return this
    }
}

And it seems to be working! However, I don’t get “Jump to source” capability for some reason:

automation_jump_to_source_missing
automation_jump_to_source_missing458×243 20.6 KB

Unwrapped elements that already worked without the target element evaluator (because they are already named element, such as YAMLKeyValue) do have “Jump to source” capability:

script_jump_to_source_present
script_jump_to_source_present514×242 23.6 KB

YAML for this case is:

script:
  bedroom_turn_on_light_conditional: # the key is the reference in this case
    sequence:
      [...]

I must be doing something wrong here, however since the TargetElementEvaluator API is not documented I did my best to understand how it worked.

Here is the complete code to the plugin: GitHub - daniele-athome/hass-intellij-plugin: Home Assistant IntelliJ plugin

Thank you!