Apply requested changes

square · Apr 15, 2022 · e6316a0 · e6316a0
1 parent f2ace6f
commit e6316a0
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 21 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,18 @@ internal class CodeWriter constructor(
     }
   }
 
+  /**
+   * Emits the `context` block for [contextReceivers].
+   */
+  fun emitContextReceivers(contextReceivers: List<TypeName>) {
+    if (contextReceivers.isNotEmpty()) {
+      val receivers = contextReceivers
+        .map { CodeBlock.of("%T", it) }
+        .joinToCode(prefix = "context(", suffix = ")")
+      emitCode(receivers)
+    }
+  }
+
   /**
    * 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/FunSpec.kt b/kotlinpoet/src/main/java/com/squareup/kotlinpoet/FunSpec.kt
@@ -88,12 +88,8 @@ public class FunSpec private constructor(
       codeWriter.emitKdoc(kdoc.ensureEndsWithNewLine())
     }
     codeWriter.emitAnnotations(annotations, false)
-    if (contextReceiverTypes.isNotEmpty()) {
-      val receivers = contextReceiverTypes
-        .map { CodeBlock.of("%T", it) }
-        .joinToCode(prefix = "context(", suffix = ")\n"))
-      codeWriter.emitCode(receivers)
-    }
+    codeWriter.emitContextReceivers(contextReceiverTypes)
+    codeWriter.emit("\n")
     codeWriter.emitModifiers(modifiers, implicitModifiers)
 
     if (!isConstructor && !name.isAccessor) {

diff --git a/kotlinpoet/src/main/java/com/squareup/kotlinpoet/LambdaTypeName.kt b/kotlinpoet/src/main/java/com/squareup/kotlinpoet/LambdaTypeName.kt
@@ -57,10 +57,8 @@ public class LambdaTypeName private constructor(
   }
 
   override fun emit(out: CodeWriter): CodeWriter {
-    if (contextReceivers.isNotEmpty()) {
-      val receivers = contextReceivers.joinToString(", ") { "%T" }
-      out.emitCode("context($receivers)·", *contextReceivers.toTypedArray())
-    }
+    out.emitContextReceivers(contextReceivers)
+    out.emit(" ")
     if (isNullable) {
       out.emit("(")
     }

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