Add support for context-receivers

kotlinpoet/src/main/java/com/squareup/kotlinpoet/ExperimentalKotlinPoetApi.kt

@@ -0,0 +1,30 @@
+/*
+ * Copyright (C) 2022 Square, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.squareup.kotlinpoet
+
+import kotlin.annotation.AnnotationTarget.CLASS
+import kotlin.annotation.AnnotationTarget.FUNCTION
+import kotlin.annotation.AnnotationTarget.PROPERTY
+import kotlin.annotation.AnnotationTarget.TYPEALIAS
+
+/**
+ * Indicates that a given API is part of an experimental KotlinPoet and is
+ * subject to API changes.
+ */
+@RequiresOptIn
+@Retention(AnnotationRetention.BINARY)
+@Target(CLASS, FUNCTION, PROPERTY, TYPEALIAS)
+public annotation class ExperimentalKotlinPoetApi

kotlinpoet/src/main/java/com/squareup/kotlinpoet/FunSpec.kt

@@ -32,6 +32,7 @@ import kotlin.DeprecationLevel.WARNING
 import kotlin.reflect.KClass
 
 /** A generated function declaration. */
+@OptIn(ExperimentalKotlinPoetApi::class)
 public class FunSpec private constructor(
   builder: Builder,
   private val tagMap: TagMap = builder.buildTagMap(),
@@ -45,6 +46,9 @@ public class FunSpec private constructor(
   public val modifiers: Set<KModifier> = builder.modifiers.toImmutableSet()
   public val typeVariables: List<TypeVariableName> = builder.typeVariables.toImmutableList()
   public val receiverType: TypeName? = builder.receiverType
+
+  @ExperimentalKotlinPoetApi
+  public val contextReceiverTypes: List<TypeName> = builder.contextReceiverTypes.toImmutableList()
   public val returnType: TypeName? = builder.returnType
   public val parameters: List<ParameterSpec> = builder.parameters.toImmutableList()
   public val delegateConstructor: String? = builder.delegateConstructor
@@ -84,6 +88,10 @@ public class FunSpec private constructor(
       codeWriter.emitKdoc(kdoc.ensureEndsWithNewLine())
     }
     codeWriter.emitAnnotations(annotations, false)
+    if (contextReceiverTypes.isNotEmpty()) {
+      val receivers = contextReceiverTypes.joinToString(", ") { "%T" }
+      codeWriter.emitCode("context($receivers)\n", *contextReceiverTypes.toTypedArray())
+    }
     codeWriter.emitModifiers(modifiers, implicitModifiers)
 
     if (!isConstructor && !name.isAccessor) {
@@ -285,6 +293,7 @@ public class FunSpec private constructor(
     internal var returnKdoc = CodeBlock.EMPTY
     internal var receiverKdoc = CodeBlock.EMPTY
     internal var receiverType: TypeName? = null
+    internal val contextReceiverTypes: MutableList<TypeName> = mutableListOf()
     internal var returnType: TypeName? = null
     internal var delegateConstructor: String? = null
     internal var delegateConstructorArguments = listOf<CodeBlock>()
@@ -359,6 +368,15 @@ public class FunSpec private constructor(
       typeVariables += typeVariable
     }
 
+    @ExperimentalKotlinPoetApi
+    public fun contextReceiver(receiverTypes: Collection<TypeName>): Builder = apply {
+      check(!name.isConstructor) { "$name: constructors cannot have context receivers" }
+      contextReceiverTypes += receiverTypes
+    }
+
+    @ExperimentalKotlinPoetApi
+    public fun contextReceiver(vararg receiverType: TypeName): Builder = contextReceiver(receiverType.toList())
+
     @JvmOverloads public fun receiver(
       receiverType: TypeName,
       kdoc: CodeBlock = CodeBlock.EMPTY

kotlinpoet/src/main/java/com/squareup/kotlinpoet/LambdaTypeName.kt

@@ -17,8 +17,11 @@ package com.squareup.kotlinpoet
 
 import kotlin.reflect.KClass
 
+@OptIn(ExperimentalKotlinPoetApi::class)
 public class LambdaTypeName private constructor(
   public val receiver: TypeName? = null,
+  @property:ExperimentalKotlinPoetApi
+  public val contextReceivers: List<TypeName> = emptyList(),
   parameters: List<ParameterSpec> = emptyList(),
   public val returnType: TypeName = UNIT,
   nullable: Boolean = false,
@@ -50,10 +53,14 @@ public class LambdaTypeName private constructor(
     suspending: Boolean = this.isSuspending,
     tags: Map<KClass<*>, Any> = this.tags.toMap()
   ): LambdaTypeName {
-    return LambdaTypeName(receiver, parameters, returnType, nullable, suspending, annotations, tags)
+    return LambdaTypeName(receiver, contextReceivers, parameters, returnType, nullable, suspending, annotations, tags)
   }
 
   override fun emit(out: CodeWriter): CodeWriter {
+    if (contextReceivers.isNotEmpty()) {
+      val receivers = contextReceivers.joinToString(", ") { "%T" }
+      out.emitCode("context($receivers)·", *contextReceivers.toTypedArray())
+    }
     if (isNullable) {
       out.emit("(")
     }
@@ -80,12 +87,20 @@ public class LambdaTypeName private constructor(
   }
 
   public companion object {
+    /** Returns a lambda type with `returnType` and parameters listed in `parameters`. */
+    @ExperimentalKotlinPoetApi @JvmStatic public fun get(
+      receiver: TypeName? = null,
+      parameters: List<ParameterSpec> = emptyList(),
+      returnType: TypeName,
+      contextReceivers: List<TypeName> = emptyList()
+    ): LambdaTypeName = LambdaTypeName(receiver, contextReceivers, parameters, returnType)
+
     /** Returns a lambda type with `returnType` and parameters listed in `parameters`. */
     @JvmStatic public fun get(
       receiver: TypeName? = null,
       parameters: List<ParameterSpec> = emptyList(),
       returnType: TypeName
-    ): LambdaTypeName = LambdaTypeName(receiver, parameters, returnType)
+    ): LambdaTypeName = LambdaTypeName(receiver, emptyList(), parameters, returnType)
 
     /** Returns a lambda type with `returnType` and parameters listed in `parameters`. */
     @JvmStatic public fun get(
@@ -95,6 +110,7 @@ public class LambdaTypeName private constructor(
     ): LambdaTypeName {
       return LambdaTypeName(
         receiver,
+        emptyList(),
         parameters.toList().map { ParameterSpec.unnamed(it) },
         returnType
       )
@@ -105,6 +121,6 @@ public class LambdaTypeName private constructor(
       receiver: TypeName? = null,
       vararg parameters: ParameterSpec = emptyArray(),
       returnType: TypeName
-    ): LambdaTypeName = LambdaTypeName(receiver, parameters.toList(), returnType)
+    ): LambdaTypeName = LambdaTypeName(receiver, emptyList(), parameters.toList(), returnType)
   }
 }

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

@@ -33,6 +33,7 @@ import javax.lang.model.util.Types
 import kotlin.test.BeforeTest
 import kotlin.test.Test
 
+@OptIn(ExperimentalKotlinPoetApi::class)
 class FunSpecTest {
   @Rule @JvmField val compilation = CompilationRule()
 
@@ -69,6 +70,8 @@ class FunSpecTest {
 
   internal interface ExtendsOthers : Callable<Int>, Comparable<Long>
 
+  annotation class TestAnnotation
+
   abstract class InvalidOverrideMethods {
     fun finalMethod() {
     }
@@ -429,6 +432,76 @@ class FunSpecTest {
     )
   }
 
+  @Test fun functionWithContextReceiver() {
+    val stringType = STRING
+    val funSpec = FunSpec.builder("foo")
+      .contextReceiver(stringType)
+      .build()
+
+    assertThat(funSpec.toString()).isEqualTo(
+      """
+      |context(kotlin.String)
+      |public fun foo(): kotlin.Unit {
+      |}
+      |""".trimMargin()
+    )
+  }
+
+  @Test fun functionWithMultipleContextReceiver() {
+    val stringType = STRING
+    val intType = INT
+    val booleanType = BOOLEAN
+    val funSpec = FunSpec.builder("foo")
+      .contextReceiver(stringType, intType, booleanType)
+      .build()
+
+    assertThat(funSpec.toString()).isEqualTo(
+      """
+      |context(kotlin.String, kotlin.Int, kotlin.Boolean)
+      |public fun foo(): kotlin.Unit {
+      |}
+      |""".trimMargin()
+    )
+  }
+
+  @Test fun functionWithGenericContextReceiver() {
+    val genericType = TypeVariableName("T")
+    val funSpec = FunSpec.builder("foo")
+      .addTypeVariable(genericType)
+      .contextReceiver(genericType)
+      .build()
+
+    assertThat(funSpec.toString()).isEqualTo(
+      """
+      |context(T)
+      |public fun <T> foo(): kotlin.Unit {
+      |}
+      |""".trimMargin()
+    )
+  }
+
+  @Test fun functionWithAnnotatedContextReceiver() {
+    val genericType = STRING.copy(annotations = listOf(AnnotationSpec.get(TestAnnotation())))
+    val funSpec = FunSpec.builder("foo")
+      .contextReceiver(genericType)
+      .build()
+
+    assertThat(funSpec.toString()).isEqualTo(
+      """
+      |context(@com.squareup.kotlinpoet.FunSpecTest.TestAnnotation kotlin.String)
+      |public fun foo(): kotlin.Unit {
+      |}
+      |""".trimMargin()
+    )
+  }
+
+  @Test fun constructorWithContextReceiver() {
+    assertThrows<IllegalStateException> {
+      FunSpec.constructorBuilder()
+        .contextReceiver(STRING)
+    }.hasMessageThat().isEqualTo("constructor(): constructors cannot have context receivers")
+  }
+
   @Test fun functionParamSingleLambdaParam() {
     val unitType = UNIT
     val booleanType = BOOLEAN

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

@@ -20,6 +20,7 @@ import com.squareup.kotlinpoet.KModifier.VARARG
 import javax.annotation.Nullable
 import kotlin.test.Test
 
+@OptIn(ExperimentalKotlinPoetApi::class)
 class LambdaTypeNameTest {
 
   @Retention(AnnotationRetention.RUNTIME)
@@ -51,6 +52,58 @@ class LambdaTypeNameTest {
     )
   }
 
+  @Test fun contextReceiver() {
+    val typeName = LambdaTypeName.get(
+      Int::class.asTypeName(),
+      listOf(),
+      Unit::class.asTypeName(),
+      listOf(STRING)
+    )
+    assertThat(typeName.toString()).isEqualTo(
+      "context(kotlin.String) kotlin.Int.() -> kotlin.Unit"
+    )
+  }
+
+  @Test fun functionWithMultipleContextReceiver() {
+    val typeName = LambdaTypeName.get(
+      Int::class.asTypeName(),
+      listOf(),
+      Unit::class.asTypeName(),
+      listOf(STRING, BOOLEAN)
+    )
+    assertThat(typeName.toString()).isEqualTo(
+      "context(kotlin.String, kotlin.Boolean) kotlin.Int.() -> kotlin.Unit"
+    )
+  }
+
+  @Test fun functionWithGenericContextReceiver() {
+    val genericType = TypeVariableName("T")
+    val typeName = LambdaTypeName.get(
+      Int::class.asTypeName(),
+      listOf(),
+      Unit::class.asTypeName(),
+      listOf(genericType)
+    )
+
+    assertThat(typeName.toString()).isEqualTo(
+      "context(T) kotlin.Int.() -> kotlin.Unit"
+    )
+  }
+
+  @Test fun functionWithAnnotatedContextReceiver() {
+    val annotatedType = STRING.copy(annotations = listOf(AnnotationSpec.get(FunSpecTest.TestAnnotation())))
+    val typeName = LambdaTypeName.get(
+      Int::class.asTypeName(),
+      listOf(),
+      Unit::class.asTypeName(),
+      listOf(annotatedType)
+    )
+
+    assertThat(typeName.toString()).isEqualTo(
+      "context(@com.squareup.kotlinpoet.FunSpecTest.TestAnnotation kotlin.String) kotlin.Int.() -> kotlin.Unit"
+    )
+  }
+
   @Test fun paramsWithAnnotationsForbidden() {
     assertThrows<IllegalArgumentException> {
       LambdaTypeName.get(