Add support for context-receivers (#1233)

* Add API for context receivers * Use existing opt-in annotation and make context-receivers not nullable * Remove collection overloads * Revert unwanted code style changes * Add wrongly remove @jvmoverloads annotation * Add code generator and tests * Apply requested changes - Use proper experimental annotation - Remove not needed overloads - Move contextReceivers parameter to the end * Update kotlinpoet/src/main/java/com/squareup/kotlinpoet/FunSpec.kt Co-authored-by: Zac Sweers <pandanomic@gmail.com> * Apply requested changes - Add more tests - Add new-line after context() * Update test for suggestion * Apply suggestions from code review Co-authored-by: Egor Andreevich <github@egorand.dev> * Apply requested changes * Rename contextReceiver to contextReceivers * Fix compiler error * Update kotlinpoet/src/main/java/com/squareup/kotlinpoet/FunSpec.kt Co-authored-by: Egor Andreevich <github@egorand.dev> * Fix compiler errors in tests * Run spotless * Update kotlinpoet/src/main/java/com/squareup/kotlinpoet/LambdaTypeName.kt Co-authored-by: Zac Sweers <pandanomic@gmail.com> * Apply requested changes * Don't emit context receiver spacing if there are no context receivers * Apply suggestions from code review Co-authored-by: Egor Andreevich <github@egorand.dev> Co-authored-by: Zac Sweers <pandanomic@gmail.com> Co-authored-by: Egor Andreevich <github@egorand.dev>
square · Apr 18, 2022 · 6d919af · 6d919af
1 parent eee9256
commit 6d919af
Show file tree

Hide file tree

Showing 6 changed files with 205 additions and 9 deletions.
diff --git a/kotlinpoet/src/main/java/com/squareup/kotlinpoet/CodeWriter.kt b/kotlinpoet/src/main/java/com/squareup/kotlinpoet/CodeWriter.kt
@@ -167,6 +167,19 @@ internal class CodeWriter constructor(
     }
   }
 
+  /**
+   * Emits the `context` block for [contextReceivers].
+   */
+  fun emitContextReceivers(contextReceivers: List<TypeName>, suffix: String = "") {
+    if (contextReceivers.isNotEmpty()) {
+      val receivers = contextReceivers
+        .map { CodeBlock.of("%T", it) }
+        .joinToCode(prefix = "context(", suffix = ")")
+      emitCode(receivers)
+      emit(suffix)
+    }
+  }
+
   /**
    * Emit type variables with their bounds. If a type variable has more than a single bound - call
    * [emitWhereBlock] with same input to produce an additional `where` block.

diff --git a/kotlinpoet/src/main/java/com/squareup/kotlinpoet/ExperimentalKotlinPoetApi.kt b/kotlinpoet/src/main/java/com/squareup/kotlinpoet/ExperimentalKotlinPoetApi.kt
@@ -0,0 +1,29 @@
+/*
+ * 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 experimental and subject to change.
+ */
+@RequiresOptIn
+@Retention(AnnotationRetention.BINARY)
+@Target(CLASS, FUNCTION, PROPERTY, TYPEALIAS)
+public annotation class ExperimentalKotlinPoetApi
diff --git a/kotlinpoet/src/main/java/com/squareup/kotlinpoet/FunSpec.kt b/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,7 @@ public class FunSpec private constructor(
       codeWriter.emitKdoc(kdoc.ensureEndsWithNewLine())
     }
     codeWriter.emitAnnotations(annotations, false)
+    codeWriter.emitContextReceivers(contextReceiverTypes, suffix = "\n")
     codeWriter.emitModifiers(modifiers, implicitModifiers)
 
     if (!isConstructor && !name.isAccessor) {
@@ -285,6 +290,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 +365,15 @@ public class FunSpec private constructor(
       typeVariables += typeVariable
     }
 
+    @ExperimentalKotlinPoetApi
+    public fun contextReceivers(receiverTypes: Iterable<TypeName>): Builder = apply {
+      check(!name.isConstructor) { "constructors cannot have context receivers" }
+      contextReceiverTypes += receiverTypes
+    }
+
+    @ExperimentalKotlinPoetApi
+    public fun contextReceivers(vararg receiverType: TypeName): Builder = contextReceivers(receiverType.toList())
+
     @JvmOverloads public fun receiver(
       receiverType: TypeName,
       kdoc: CodeBlock = CodeBlock.EMPTY

diff --git a/kotlinpoet/src/main/java/com/squareup/kotlinpoet/LambdaTypeName.kt b/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,11 @@ 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 {
+    out.emitContextReceivers(contextReceivers, suffix = "·")
     if (isNullable) {
       out.emit("(")
     }
@@ -80,12 +84,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 +107,7 @@ public class LambdaTypeName private constructor(
     ): LambdaTypeName {
       return LambdaTypeName(
         receiver,
+        emptyList(),
         parameters.toList().map { ParameterSpec.unnamed(it) },
         returnType
       )
@@ -105,6 +118,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)
   }
 }
diff --git a/kotlinpoet/src/test/java/com/squareup/kotlinpoet/FunSpecTest.kt b/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")
+      .contextReceivers(stringType)
+      .build()
+
+    assertThat(funSpec.toString()).isEqualTo(
+      """
+      |context(kotlin.String)
+      |public fun foo(): kotlin.Unit {
+      |}
+      |""".trimMargin()
+    )
+  }
+
+  @Test fun functionWithMultipleContextReceivers() {
+    val stringType = STRING
+    val intType = INT
+    val booleanType = BOOLEAN
+    val funSpec = FunSpec.builder("foo")
+      .contextReceivers(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)
+      .contextReceivers(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")
+      .contextReceivers(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()
+        .contextReceivers(STRING)
+    }.hasMessageThat().isEqualTo("constructors cannot have context receivers")
+  }
+
   @Test fun functionParamSingleLambdaParam() {
     val unitType = UNIT
     val booleanType = BOOLEAN

diff --git a/kotlinpoet/src/test/java/com/squareup/kotlinpoet/LambdaTypeNameTest.kt b/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)
@@ -30,27 +31,79 @@ class LambdaTypeNameTest {
 
   @Test fun receiverWithoutAnnotationHasNoParens() {
     val typeName = LambdaTypeName.get(
-      Int::class.asClassName(),
-      listOf(),
-      Unit::class.asTypeName()
+      receiver = Int::class.asClassName(),
+      parameters = listOf(),
+      returnType = Unit::class.asTypeName()
     )
     assertThat(typeName.toString()).isEqualTo("kotlin.Int.() -> kotlin.Unit")
   }
 
   @Test fun receiverWithAnnotationHasParens() {
     val annotation = IsAnnotated::class.java.getAnnotation(HasSomeAnnotation::class.java)
     val typeName = LambdaTypeName.get(
-      Int::class.asClassName().copy(
+      receiver = Int::class.asClassName().copy(
         annotations = listOf(AnnotationSpec.get(annotation, includeDefaultValues = true))
       ),
-      listOf(),
-      Unit::class.asTypeName()
+      parameters = listOf(),
+      returnType = Unit::class.asTypeName()
     )
     assertThat(typeName.toString()).isEqualTo(
       "(@com.squareup.kotlinpoet.LambdaTypeNameTest.HasSomeAnnotation kotlin.Int).() -> kotlin.Unit"
     )
   }
 
+  @Test fun contextReceiver() {
+    val typeName = LambdaTypeName.get(
+      receiver = Int::class.asTypeName(),
+      parameters = listOf(),
+      returnType = Unit::class.asTypeName(),
+      contextReceivers = listOf(STRING)
+    )
+    assertThat(typeName.toString()).isEqualTo(
+      "context(kotlin.String) kotlin.Int.() -> kotlin.Unit"
+    )
+  }
+
+  @Test fun functionWithMultipleContextReceivers() {
+    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(