Add support for anonymous functions

square · Oct 14, 2018 · bd06fc1 · bd06fc1
1 parent c348f8e
commit bd06fc1
Show file tree

Hide file tree

Showing 5 changed files with 42 additions and 11 deletions.
diff --git a/src/main/java/com/squareup/kotlinpoet/CodeBlock.kt b/src/main/java/com/squareup/kotlinpoet/CodeBlock.kt
@@ -333,7 +333,7 @@ class CodeBlock private constructor(
       is CharSequence -> o.toString()
       is ParameterSpec -> o.name
       is PropertySpec -> o.name
-      is FunSpec -> o.name
+      is FunSpec -> o.name!!
       is TypeSpec -> o.name!!
       else -> throw IllegalArgumentException("expected name but was " + o)
     }

diff --git a/src/main/java/com/squareup/kotlinpoet/FunSpec.kt b/src/main/java/com/squareup/kotlinpoet/FunSpec.kt
@@ -71,7 +71,9 @@ class FunSpec private constructor(builder: Builder) {
     codeWriter.emitAnnotations(annotations, false)
     codeWriter.emitModifiers(modifiers, implicitModifiers)
 
-    if (!isConstructor && !name.isAccessor) {
+    if (name == null) {
+      codeWriter.emit("fun")
+    } else if (!isConstructor && !name.isAccessor) {
       codeWriter.emit("fun ")
     }
 
@@ -117,7 +119,9 @@ class FunSpec private constructor(builder: Builder) {
           codeWriter.emitCode("%T.", receiverType)
         }
       }
-      codeWriter.emitCode("%L", escapeIfNecessary(name))
+      if (name != null) {
+        codeWriter.emitCode("%L", escapeIfNecessary(name))
+      }
     }
 
     parameters.emit(codeWriter) { param ->
@@ -181,7 +185,7 @@ class FunSpec private constructor(builder: Builder) {
     return builder
   }
 
-  class Builder internal constructor(internal val name: String) {
+  class Builder internal constructor(internal val name: String?) {
     internal val kdoc = CodeBlock.builder()
     internal var receiverType: TypeName? = null
     internal var returnType: TypeName? = null
@@ -366,8 +370,8 @@ class FunSpec private constructor(builder: Builder) {
     internal const val GETTER = "get()"
     internal const val SETTER = "set()"
 
-    internal val String.isConstructor get() = this == CONSTRUCTOR
-    internal val String.isAccessor get() = this.isOneOf(GETTER, SETTER)
+    internal val String?.isConstructor get() = this === CONSTRUCTOR
+    internal val String?.isAccessor get() = this != null && this.isOneOf(GETTER, SETTER)
 
     private val EXPRESSION_BODY_PREFIX = CodeBlock.of("return ")
 
@@ -379,6 +383,8 @@ class FunSpec private constructor(builder: Builder) {
 
     @JvmStatic fun setterBuilder() = Builder(SETTER)
 
+    @JvmStatic fun anonymousFunctionBuilder() = Builder(null)
+
     /**
      * Returns a new fun spec builder that overrides `method`.
 

diff --git a/src/main/java/com/squareup/kotlinpoet/ParameterSpec.kt b/src/main/java/com/squareup/kotlinpoet/ParameterSpec.kt
@@ -34,8 +34,9 @@ class ParameterSpec private constructor(builder: ParameterSpec.Builder) {
     codeWriter.emitAnnotations(annotations, true)
     codeWriter.emitModifiers(modifiers)
     if (name.isNotEmpty()) codeWriter.emitCode("%L", escapeIfNecessary(name))
-    if (name.isNotEmpty() && includeType) codeWriter.emit(": ")
-    if (includeType) codeWriter.emitCode("%T", type)
+    val emitType = type != null && includeType
+    if (name.isNotEmpty() && emitType) codeWriter.emit(": ")
+    if (emitType) codeWriter.emitCode("%T", type)
     emitDefaultValue(codeWriter)
   }
 
@@ -56,7 +57,7 @@ class ParameterSpec private constructor(builder: ParameterSpec.Builder) {
 
   override fun toString() = buildString { emit(CodeWriter(this)) }
 
-  fun toBuilder(name: String = this.name, type: TypeName = this.type): Builder {
+  fun toBuilder(name: String = this.name, type: TypeName? = this.type): Builder {
     val builder = Builder(name, type)
     builder.annotations += annotations
     builder.modifiers += modifiers
@@ -66,7 +67,7 @@ class ParameterSpec private constructor(builder: ParameterSpec.Builder) {
 
   class Builder internal constructor(
     internal val name: String,
-    internal val type: TypeName
+    internal val type: TypeName?
   ) {
     internal var defaultValue: CodeBlock? = null
 
@@ -152,6 +153,8 @@ class ParameterSpec private constructor(builder: ParameterSpec.Builder) {
     @JvmStatic fun unnamed(type: Type) = unnamed(type.asTypeName())
 
     @JvmStatic fun unnamed(type: TypeName) = Builder("", type).build()
+
+    @JvmStatic fun untyped(name: String) = Builder(name, null).build()
   }
 }
 

diff --git a/src/main/java/com/squareup/kotlinpoet/Util.kt b/src/main/java/com/squareup/kotlinpoet/Util.kt
@@ -55,7 +55,7 @@ internal fun requireNoneOf(modifiers: Set<KModifier>, vararg forbidden: KModifie
   }
 }
 
-internal fun <T> T.isOneOf(t1: T, t2: T, t3: T? = null, t4: T? = null, t5: T? = null, t6: T? = null) =
+internal fun <T : Any> T.isOneOf(t1: T, t2: T, t3: T? = null, t4: T? = null, t5: T? = null, t6: T? = null) =
     this == t1 || this == t2 || this == t3 || this == t4 || this == t5 || this == t6
 
 internal fun <T> Collection<T>.containsAnyOf(vararg t: T) = t.any(this::contains)

diff --git a/src/test/java/com/squareup/kotlinpoet/FunSpecTest.kt b/src/test/java/com/squareup/kotlinpoet/FunSpecTest.kt
@@ -479,4 +479,26 @@ class FunSpecTest {
 
     assertThat(builder.build().parameters).containsExactly(seasoning)
   }
+
+  @Test fun anonymousFunction() {
+    val funSpec = FunSpec.anonymousFunctionBuilder()
+        .addParameter("x", Int::class)
+        .addParameter("y", Int::class)
+        .returns(Int::class)
+        .addStatement("return x + y")
+        .build()
+    assertThat(funSpec.toString()).isEqualTo("""
+      |fun(x: kotlin.Int, y: kotlin.Int): kotlin.Int = x + y
+      |""".trimMargin())
+  }
+
+  @Test fun anonymousFunctionWithUntypedParameters() {
+    val funSpec = FunSpec.anonymousFunctionBuilder()
+        .addParameter(ParameterSpec.untyped("item"))
+        .addStatement("return item > 0")
+        .build()
+    assertThat(funSpec.toString()).isEqualTo("""
+      |fun(item) = item > 0
+      |""".trimMargin())
+  }
 }