PsiAugmentProvider for Kotlin class?

How does one use a PsiAugmentProvider to augment a Kotlin class?

 

I tried the following but it does not find the method when writing eg: this.myMethod() causes [UNRESOLVED_REFERENCE]

override fun <Psi : PsiElement?> getAugments(element: PsiElement, type: Class<Psi>): MutableList<Psi> {
if (type == PsiMethod::class.java)
if (element is KtLightClass && element.hasAnnotation("com.example.MyAnnotation")) {
val method = LightMethodBuilder(element.manager, JavaLanguage.INSTANCE, "myMethod")
.apply {
containingClass = element
setMethodReturnType("Unit")
setModifiers("public")
}

return mutableListOf(KtLightMethodImpl.create(method, null, element) as Psi)
}

return super.getAugments(element, type) as MutableList<Psi>
}

Please sign in to leave a comment.