Add support for context receivers @PropertySpec and fix issues with annotations

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
     }
 

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(

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

kotlinpoet/src/test/java/com/squareup/kotlinpoet/PropertySpecTest.kt

@@ -16,14 +16,19 @@
 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
 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()
+    )
+  }
 }