Handle multi-annotations with use-site targets

Summary: This fixes #48 Reviewed By: strulovich Differential Revision: D22101948 fbshipit-source-id: ef79e00cdabd9663cd6282abaf77649c1b976569
facebook · Jul 7, 2020 · bab0f4b · bab0f4b
1 parent d4cae4b
commit bab0f4b
Show file tree

Hide file tree

Showing 2 changed files with 52 additions and 4 deletions.
diff --git a/core/src/main/java/com/facebook/ktfmt/KotlinInputAstVisitor.kt b/core/src/main/java/com/facebook/ktfmt/KotlinInputAstVisitor.kt
@@ -35,6 +35,7 @@ import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.KtAnnotatedExpression
 import org.jetbrains.kotlin.psi.KtAnnotation
 import org.jetbrains.kotlin.psi.KtAnnotationEntry
+import org.jetbrains.kotlin.psi.KtAnnotationUseSiteTarget
 import org.jetbrains.kotlin.psi.KtArrayAccessExpression
 import org.jetbrains.kotlin.psi.KtBinaryExpression
 import org.jetbrains.kotlin.psi.KtBinaryExpressionWithTypeRHS
@@ -1159,13 +1160,42 @@ class KotlinInputAstVisitor(
     }
   }
 
+  /** For example, @field:[Inject Named("WEB_VIEW")] */
+  override fun visitAnnotation(annotation: KtAnnotation) {
+    builder.sync(annotation)
+    builder.block(ZERO) {
+      builder.token("@")
+      annotation.useSiteTarget?.accept(this)
+      builder.token(":")
+      builder.block(expressionBreakIndent) {
+        builder.token("[")
+        forEachCommaSeparated(
+            annotation.entries,
+            delimiter = { builder.breakOp(Doc.FillMode.INDEPENDENT, " ", ZERO) },
+            function = { it.accept(this) })
+      }
+      builder.token("]")
+    }
+    builder.forcedBreak()
+  }
+
+  /** For example, 'field' in @field:[Inject Named("WEB_VIEW")] */
+  override fun visitAnnotationUseSiteTarget(
+      annotationTarget: KtAnnotationUseSiteTarget, data: Void?
+  ): Void? {
+    builder.token(annotationTarget.getAnnotationUseSiteTarget().renderName)
+    return null
+  }
+
   /** For example `@Magic` or `@Fred(1, 5)` */
   override fun visitAnnotationEntry(annotationEntry: KtAnnotationEntry) {
     builder.sync(annotationEntry)
-    builder.token("@")
-    val useSiteTarget = annotationEntry.useSiteTarget?.getAnnotationUseSiteTarget()
-    if (useSiteTarget != null) {
-      builder.token(useSiteTarget.renderName)
+    if (annotationEntry.atSymbol != null) {
+      builder.token("@")
+    }
+    val useSiteTarget = annotationEntry.useSiteTarget
+    if (useSiteTarget != null && useSiteTarget.parent == annotationEntry) {
+      useSiteTarget.accept(this)
       builder.token(":")
     }
     visitCallElement(

diff --git a/core/src/test/java/com/facebook/ktfmt/FormatterKtTest.kt b/core/src/test/java/com/facebook/ktfmt/FormatterKtTest.kt
@@ -2084,6 +2084,24 @@ class FormatterKtTest {
       |""".trimMargin(),
           deduceMaxWidth = true)
 
+  @Test
+  fun `handle multi-annotations with use-site targets`() =
+      assertFormatted(
+          """
+    |class Something {
+    |  @field:[Inject Named("WEB_VIEW")]
+    |  internal lateinit var httpClient: OkHttpClient
+    |
+    |  @field:[Inject Named("WEB_VIEW")]
+    |  var httpClient: OkHttpClient
+    |
+    |  @Px
+    |  @field:[Inject Named("WEB_VIEW")]
+    |  var httpClient: OkHttpClient
+    |}
+    |
+    """.trimMargin())
+
   @Test
   fun `handle lambda types`() =
       assertFormatted(