Add support for context receivers @PropertySpec and fix issues with annotations (#1247)

* Fix context receiver and annotation order @FunSpec
Prevent context receivers on accessors
Add FunSpec tests:
 - Annotated function with context receiver
 - Accessor with context receiver
Add context receivers to PropertySpec
Add PropertySpec tests:
 - Var with context receiver
 - Val without getter with context receiver
 - Val with context receiver
 - Annotated val with context receiver

* Fix code style

* Fix checks on var with context receivers and custom accessors
Update tests for vars with context receivers without custom accessors
Add test for var with context receivers and custom accessors
Fix code style

* Update kotlinpoet/src/main/java/com/squareup/kotlinpoet/PropertySpec.kt

Co-authored-by: Egor Andreevich <github@egorand.dev>

* Update tests

Co-authored-by: Egor Andreevich <github@egorand.dev>

4 files changed, 174 insertions, 1 deletions

diff --git a/kotlinpoet/src/main/java/com/squareup/kotlinpoet/FunSpec.kt b/kotlinpoet/src/main/java/com/squareup/kotlinpoet/FunSpec.kt
index 4d5f2428..2266f6dd 100644
--- a/kotlinpoet/src/main/java/com/squareup/kotlinpoet/FunSpec.kt
+++ b/kotlinpoet/src/main/java/com/squareup/kotlinpoet/FunSpec.kt
@@ -87,8 +87,8 @@ public class FunSpec private constructor(
     } else {
       codeWriter.emitKdoc(kdoc.ensureEndsWithNewLine())
     }
-    codeWriter.emitAnnotations(annotations, false)
     codeWriter.emitContextReceivers(contextReceiverTypes, suffix = "\n")
+    codeWriter.emitAnnotations(annotations, false)
     codeWriter.emitModifiers(modifiers, implicitModifiers)
 
     if (!isConstructor && !name.isAccessor) {
@@ -368,6 +368,7 @@ public class FunSpec private constructor(
     @ExperimentalKotlinPoetApi
     public fun contextReceivers(receiverTypes: Iterable<TypeName>): Builder = apply {
       check(!name.isConstructor) { "constructors cannot have context receivers" }
+      check(!name.isAccessor) { "$name cannot have context receivers" }
       contextReceiverTypes += receiverTypes
     }
 
diff --git a/kotlinpoet/src/main/java/com/squareup/kotlinpoet/PropertySpec.kt b/kotlinpoet/src/main/java/com/squareup/kotlinpoet/PropertySpec.kt
index 9d171154..b69702f8 100644
--- a/kotlinpoet/src/main/java/com/squareup/kotlinpoet/PropertySpec.kt
+++ b/kotlinpoet/src/main/java/com/squareup/kotlinpoet/PropertySpec.kt
@@ -24,6 +24,7 @@ import javax.lang.model.element.Element
 import kotlin.reflect.KClass
 
 /** A generated property declaration. */
+@OptIn(ExperimentalKotlinPoetApi::class)
 public class PropertySpec private constructor(
   builder: Builder,
   private val tagMap: TagMap = builder.buildTagMap(),
@@ -42,6 +43,9 @@ public class PropertySpec private constructor(
   public val setter: FunSpec? = builder.setter
   public val receiverType: TypeName? = builder.receiverType
 
+  @ExperimentalKotlinPoetApi
+  public val contextReceiverTypes: List<TypeName> = builder.contextReceiverTypes.toImmutableList()
+
   init {
     require(
       typeVariables.none { it.isReified } ||
@@ -54,6 +58,12 @@ public class PropertySpec private constructor(
     require(mutable || setter == null) {
       "only a mutable property can have a setter"
     }
+    if (contextReceiverTypes.isNotEmpty()) {
+      requireNotNull(getter) { "properties with context receivers require a $GETTER" }
+      if (mutable) {
+        requireNotNull(setter) { "mutable properties with context receivers require a $SETTER" }
+      }
+    }
   }
 
   internal fun emit(
@@ -70,6 +80,7 @@ public class PropertySpec private constructor(
     if (emitKdoc) {
       codeWriter.emitKdoc(kdoc.ensureEndsWithNewLine())
     }
+    codeWriter.emitContextReceivers(contextReceiverTypes, suffix = "\n")
     codeWriter.emitAnnotations(annotations, inlineAnnotations)
     codeWriter.emitModifiers(propertyModifiers, implicitModifiers)
     codeWriter.emitCode(if (mutable) "var·" else "val·")
@@ -174,6 +185,7 @@ public class PropertySpec private constructor(
     internal var getter: FunSpec? = null
     internal var setter: FunSpec? = null
     internal var receiverType: TypeName? = null
+    internal val contextReceiverTypes: MutableList<TypeName> = mutableListOf()
 
     public val annotations: MutableList<AnnotationSpec> = mutableListOf()
     public val modifiers: MutableList<KModifier> = mutableListOf()
@@ -270,6 +282,15 @@ public class PropertySpec private constructor(
 
     public fun receiver(receiverType: KClass<*>): Builder = receiver(receiverType.asTypeName())
 
+    @ExperimentalKotlinPoetApi
+    public fun contextReceivers(receiverTypes: Iterable<TypeName>): Builder = apply {
+      contextReceiverTypes += receiverTypes
+    }
+
+    @ExperimentalKotlinPoetApi
+    public fun contextReceivers(vararg receiverType: TypeName): Builder =
+      contextReceivers(receiverType.toList())
+
     public fun build(): PropertySpec {
       if (KModifier.INLINE in modifiers) {
         throw IllegalArgumentException(
diff --git a/kotlinpoet/src/test/java/com/squareup/kotlinpoet/FunSpecTest.kt b/kotlinpoet/src/test/java/com/squareup/kotlinpoet/FunSpecTest.kt
index 8804e7b3..df93cfcb 100644
--- a/kotlinpoet/src/test/java/com/squareup/kotlinpoet/FunSpecTest.kt
+++ b/kotlinpoet/src/test/java/com/squareup/kotlinpoet/FunSpecTest.kt
@@ -18,6 +18,8 @@ package com.squareup.kotlinpoet
 import com.google.common.collect.Iterables.getOnlyElement
 import com.google.common.truth.Truth.assertThat
 import com.google.testing.compile.CompilationRule
+import com.squareup.kotlinpoet.FunSpec.Companion.GETTER
+import com.squareup.kotlinpoet.FunSpec.Companion.SETTER
 import com.squareup.kotlinpoet.ParameterizedTypeName.Companion.parameterizedBy
 import org.junit.Rule
 import java.io.Closeable
@@ -480,6 +482,22 @@ class FunSpecTest {
     )
   }
 
+  @Test fun annotatedFunctionWithContextReceiver() {
+    val funSpec = FunSpec.builder("foo")
+      .addAnnotation(AnnotationSpec.get(TestAnnotation()))
+      .contextReceivers(STRING)
+      .build()
+
+    assertThat(funSpec.toString()).isEqualTo(
+      """
+      |context(kotlin.String)
+      |@com.squareup.kotlinpoet.FunSpecTest.TestAnnotation
+      |public fun foo(): kotlin.Unit {
+      |}
+      |""".trimMargin()
+    )
+  }
+
   @Test fun functionWithAnnotatedContextReceiver() {
     val genericType = STRING.copy(annotations = listOf(AnnotationSpec.get(TestAnnotation())))
     val funSpec = FunSpec.builder("foo")
@@ -502,6 +520,18 @@ class FunSpecTest {
     }.hasMessageThat().isEqualTo("constructors cannot have context receivers")
   }
 
+  @Test fun accessorWithContextReceiver() {
+    assertThrows<IllegalStateException> {
+      FunSpec.getterBuilder()
+        .contextReceivers(STRING)
+    }.hasMessageThat().isEqualTo("$GETTER cannot have context receivers")
+
+    assertThrows<IllegalStateException> {
+      FunSpec.setterBuilder()
+        .contextReceivers(STRING)
+    }.hasMessageThat().isEqualTo("$SETTER cannot have context receivers")
+  }
+
   @Test fun functionParamSingleLambdaParam() {
     val unitType = UNIT
     val booleanType = BOOLEAN
diff --git a/kotlinpoet/src/test/java/com/squareup/kotlinpoet/PropertySpecTest.kt b/kotlinpoet/src/test/java/com/squareup/kotlinpoet/PropertySpecTest.kt
index 2ab1e244..ca80ba44 100644
--- a/kotlinpoet/src/test/java/com/squareup/kotlinpoet/PropertySpecTest.kt
+++ b/kotlinpoet/src/test/java/com/squareup/kotlinpoet/PropertySpecTest.kt
@@ -16,6 +16,8 @@
 package com.squareup.kotlinpoet
 
 import com.google.common.truth.Truth.assertThat
+import com.squareup.kotlinpoet.FunSpec.Companion.GETTER
+import com.squareup.kotlinpoet.FunSpec.Companion.SETTER
 import com.squareup.kotlinpoet.KModifier.PRIVATE
 import com.squareup.kotlinpoet.ParameterizedTypeName.Companion.parameterizedBy
 import java.io.Serializable
@@ -23,7 +25,10 @@ import java.util.function.Function
 import kotlin.reflect.KClass
 import kotlin.test.Test
 
+@OptIn(ExperimentalKotlinPoetApi::class)
 class PropertySpecTest {
+  annotation class TestAnnotation
+
   @Test fun nullable() {
     val type = String::class.asClassName().copy(nullable = true)
     val a = PropertySpec.builder("foo", type).build()
@@ -559,4 +564,120 @@ class PropertySpecTest {
       """.trimIndent()
     )
   }
+
+  @Test fun varWithContextReceiverWithoutCustomAccessors() {
+    val mutablePropertySpecBuilder = {
+      PropertySpec.builder("foo", STRING)
+        .mutable()
+        .contextReceivers(INT)
+    }
+
+    assertThrows<IllegalArgumentException> {
+      mutablePropertySpecBuilder()
+        .getter(
+          FunSpec.getterBuilder()
+            .build()
+        )
+        .build()
+    }.hasMessageThat()
+      .isEqualTo("mutable properties with context receivers require a $SETTER")
+
+    assertThrows<IllegalArgumentException> {
+      mutablePropertySpecBuilder()
+        .setter(
+          FunSpec.setterBuilder()
+            .build()
+        )
+        .build()
+    }.hasMessageThat()
+      .isEqualTo("properties with context receivers require a $GETTER")
+  }
+
+  @Test fun valWithContextReceiverWithoutGetter() {
+    assertThrows<IllegalArgumentException> {
+      PropertySpec.builder("foo", STRING)
+        .mutable(false)
+        .contextReceivers(INT)
+        .build()
+    }.hasMessageThat()
+      .isEqualTo("properties with context receivers require a $GETTER")
+  }
+
+  @Test fun varWithContextReceiver() {
+    val propertySpec = PropertySpec.builder("foo", INT)
+      .mutable()
+      .contextReceivers(STRING)
+      .getter(
+        FunSpec.getterBuilder()
+          .addStatement("return \"\"")
+          .build()
+      )
+      .setter(
+        FunSpec.setterBuilder()
+          .addParameter(
+            ParameterSpec.builder("value", STRING)
+              .build()
+          )
+          .addStatement("")
+          .build()
+      )
+      .build()
+
+    assertThat(propertySpec.toString()).isEqualTo(
+      """
+      |context(kotlin.String)
+      |var foo: kotlin.Int
+      |  get() = ""
+      |  set(`value`) {
+      |
+      |  }
+      |
+      """.trimMargin()
+    )
+  }
+
+  @Test fun valWithContextReceiver() {
+    val propertySpec = PropertySpec.builder("foo", INT)
+      .mutable(false)
+      .contextReceivers(STRING)
+      .getter(
+        FunSpec.getterBuilder()
+          .addStatement("return length")
+          .build()
+      )
+      .build()
+
+    assertThat(propertySpec.toString()).isEqualTo(
+      """
+      |context(kotlin.String)
+      |val foo: kotlin.Int
+      |  get() = length
+      |
+      """.trimMargin()
+    )
+  }
+
+  @OptIn(DelicateKotlinPoetApi::class)
+  @Test fun annotatedValWithContextReceiver() {
+    val propertySpec = PropertySpec.builder("foo", INT)
+      .mutable(false)
+      .addAnnotation(AnnotationSpec.get(TestAnnotation()))
+      .contextReceivers(STRING)
+      .getter(
+        FunSpec.getterBuilder()
+          .addStatement("return length")
+          .build()
+      )
+      .build()
+
+    assertThat(propertySpec.toString()).isEqualTo(
+      """
+      |context(kotlin.String)
+      |@com.squareup.kotlinpoet.PropertySpecTest.TestAnnotation
+      |val foo: kotlin.Int
+      |  get() = length
+      |
+      """.trimMargin()
+    )
+  }
 }