Repository: square/anvil
Pull request #470 (2022-01-27 08:07:35)

Support custom code generators through the `AnvilCompilation` API in unit tests.
I suggest only reviewing the 2nd commit. The first commit is only moving code into a separate file.

This is pretty helpful for #462. Sample usage will look like that:
```kotlin
  @Test fun `an @ContributeSubcomponent class can be generated`() {
    val codeGenerator = object : CodeGenerator {
      override fun isApplicable(context: AnvilContext): Boolean = true

      override fun generateCode(
        codeGenDir: File,
        module: ModuleDescriptor,
        projectFiles: Collection<KtFile>
      ): Collection<GeneratedFile> {
        return ...
      }
    }

    compile(
      """
        code
      """,
      codeGenerators = listOf(codeGenerator)
    ) {
       ...
    }
  }
```


Patch:
diff --git a/compiler/src/main/java/com/squareup/anvil/compiler/AnvilComponentRegistrar.kt b/compiler/src/main/java/com/squareup/anvil/compiler/AnvilComponentRegistrar.kt
index 82960f6a..0f0f400c 100644
--- a/compiler/src/main/java/com/squareup/anvil/compiler/AnvilComponentRegistrar.kt
+++ b/compiler/src/main/java/com/squareup/anvil/compiler/AnvilComponentRegistrar.kt
@@ -1,6 +1,7 @@
 package com.squareup.anvil.compiler
 
 import com.google.auto.service.AutoService
+import com.squareup.anvil.annotations.ExperimentalAnvilApi
 import com.squareup.anvil.compiler.api.CodeGenerator
 import com.squareup.anvil.compiler.codegen.BindingModuleGenerator
 import com.squareup.anvil.compiler.codegen.CodeGenerationExtension
@@ -23,6 +24,9 @@ import java.util.ServiceLoader
  */
 @AutoService(ComponentRegistrar::class)
 class AnvilComponentRegistrar : ComponentRegistrar {
+
+  private val manuallyAddedCodeGenerators = mutableListOf<CodeGenerator>()
+
   override fun registerProjectComponents(
     project: MockProject,
     configuration: CompilerConfiguration
@@ -39,6 +43,7 @@ class AnvilComponentRegistrar : ComponentRegistrar {
       disableComponentMerging
     )
     val codeGenerators = loadCodeGenerators()
+      .plus(manuallyAddedCodeGenerators)
       .filter { it.isApplicable(context) }
       .toMutableList()
 
@@ -94,4 +99,10 @@ class AnvilComponentRegistrar : ComponentRegistrar {
     return ServiceLoader.load(CodeGenerator::class.java, CodeGenerator::class.java.classLoader)
       .toList()
   }
+
+  // This function is used in tests. It must be called before the test code is being compiled.
+  @ExperimentalAnvilApi
+  fun addCodeGenerators(codeGenerators: List<CodeGenerator>) {
+    manuallyAddedCodeGenerators += codeGenerators
+  }
 }


Test patch:
diff --git a/compiler-utils/src/testFixtures/java/com/squareup/anvil/compiler/internal/testing/AnvilCompilation.kt b/compiler-utils/src/testFixtures/java/com/squareup/anvil/compiler/internal/testing/AnvilCompilation.kt
new file mode 100644
index 00000000..4e40eb2e
--- /dev/null
+++ b/compiler-utils/src/testFixtures/java/com/squareup/anvil/compiler/internal/testing/AnvilCompilation.kt
@@ -0,0 +1,260 @@
+package com.squareup.anvil.compiler.internal.testing
+
+import com.squareup.anvil.annotations.ExperimentalAnvilApi
+import com.squareup.anvil.compiler.AnvilCommandLineProcessor
+import com.squareup.anvil.compiler.AnvilComponentRegistrar
+import com.squareup.anvil.compiler.api.CodeGenerator
+import com.tschuchort.compiletesting.KotlinCompilation
+import com.tschuchort.compiletesting.KotlinCompilation.Result
+import com.tschuchort.compiletesting.PluginOption
+import com.tschuchort.compiletesting.SourceFile
+import dagger.internal.codegen.ComponentProcessor
+import org.intellij.lang.annotations.Language
+import org.jetbrains.kotlin.config.JvmTarget
+import java.io.File
+import java.io.OutputStream
+import java.net.URLClassLoader
+import java.nio.file.Files
+
+/**
+ * A simple API over a [KotlinCompilation] with extra configuration support for Anvil.
+ */
+@ExperimentalAnvilApi
+public class AnvilCompilation internal constructor(
+  val kotlinCompilation: KotlinCompilation
+) {
+
+  private var isCompiled = false
+  private var anvilConfigured = false
+  private var swapClassLoader = false
+
+  /** Configures this the Anvil behavior of this compilation. */
+  @Suppress("SuspiciousCollectionReassignment")
+  @ExperimentalAnvilApi
+  public fun configureAnvil(
+    enableDaggerAnnotationProcessor: Boolean = false,
+    generateDaggerFactories: Boolean = false,
+    generateDaggerFactoriesOnly: Boolean = false,
+    disableComponentMerging: Boolean = false,
+    enableExperimentalAnvilApis: Boolean = true,
+    codeGenerators: List<CodeGenerator> = emptyList(),
+  ) = apply {
+    checkNotCompiled()
+    anvilConfigured = true
+    kotlinCompilation.apply {
+      compilerPlugins = listOf(
+        AnvilComponentRegistrar().also { it.addCodeGenerators(codeGenerators) }
+      )
+      if (enableDaggerAnnotationProcessor) {
+        annotationProcessors = listOf(ComponentProcessor())
+      }
+
+      val anvilCommandLineProcessor = AnvilCommandLineProcessor()
+      commandLineProcessors = listOf(anvilCommandLineProcessor)
+
+      pluginOptions = listOf(
+        PluginOption(
+          pluginId = anvilCommandLineProcessor.pluginId,
+          optionName = "src-gen-dir",
+          optionValue = File(workingDir, "build/anvil").absolutePath
+        ),
+        PluginOption(
+          pluginId = anvilCommandLineProcessor.pluginId,
+          optionName = "generate-dagger-factories",
+          optionValue = generateDaggerFactories.toString()
+        ),
+        PluginOption(
+          pluginId = anvilCommandLineProcessor.pluginId,
+          optionName = "generate-dagger-factories-only",
+          optionValue = generateDaggerFactoriesOnly.toString()
+        ),
+        PluginOption(
+          pluginId = anvilCommandLineProcessor.pluginId,
+          optionName = "disable-component-merging",
+          optionValue = disableComponentMerging.toString()
+        )
+      )
+
+      if (enableExperimentalAnvilApis) {
+        kotlincArguments += listOf(
+          "-Xopt-in=kotlin.RequiresOptIn",
+          "-Xopt-in=com.squareup.anvil.annotations.ExperimentalAnvilApi"
+        )
+      }
+    }
+  }
+
+  /**
+   * Helpful shim to configure both [KotlinCompilation.useIR] and [KotlinCompilation.useOldBackend]
+   * accordingly.
+   */
+  public fun useIR(useIR: Boolean) = apply {
+    checkNotCompiled()
+    kotlinCompilation.useIR = useIR
+    kotlinCompilation.useOldBackend = !useIR
+
+    if (!useIR) {
+      // To suppress:
+      //  w: Language version is automatically inferred to 1.5 when using the old JVM backend.
+      //    Consider specifying -language-version explicitly, or remove -Xuse-old-backend
+      //  e: warnings found and -Werror specified
+      kotlinCompilation.languageVersion = "1.5"
+
+      // To suppress:
+      //  w: -Xuse-old-backend is deprecated and will be removed in a future release
+      //  e: warnings found and -Werror specified
+      kotlinCompilation.allWarningsAsErrors = false
+    }
+  }
+
+  /** Adds the given sources to this compilation with their packages and names inferred. */
+  public fun addSources(@Language("kotlin") vararg sources: String) = apply {
+    checkNotCompiled()
+    kotlinCompilation.sources += sources.mapIndexed { index, content ->
+      val packageDir = content.lines()
+        .firstOrNull { it.trim().startsWith("package ") }
+        ?.substringAfter("package ")
+        ?.replace('.', '/')
+        ?.let { "$it/" }
+        ?: ""
+
+      val name = "${kotlinCompilation.workingDir.absolutePath}/sources/src/main/java/" +
+        "$packageDir/Source$index.kt"
+
+      Files.createDirectories(File(name).parentFile.toPath())
+
+      SourceFile.kotlin(name, contents = content, trimIndent = true)
+    }
+  }
+
+  public fun addPreviousCompilationResult(result: Result) = apply {
+    checkNotCompiled()
+    kotlinCompilation.classpaths += result.outputDirectory
+    swapClassLoader = true
+  }
+
+  /**
+   * Returns an Anvil-generated file with the given [packageName] and [fileName] from its expected
+   * path.
+   */
+  public fun generatedAnvilFile(
+    packageName: String,
+    fileName: String
+  ): File {
+    check(isCompiled) {
+      "No compilation run yet! Call compile() first."
+    }
+    return File(
+      kotlinCompilation.workingDir,
+      "build/anvil/${packageName.replace('.', File.separatorChar)}/$fileName.kt"
+    )
+      .apply {
+        check(exists()) {
+          "Generated file not found!"
+        }
+      }
+  }
+
+  private fun checkNotCompiled() {
+    check(!isCompiled) {
+      "Already compiled! Create a new compilation if you want to compile again."
+    }
+  }
+
+  /**
+   * Compiles the underlying [KotlinCompilation]. Note that if [configureAnvil] has not been called
+   * prior to this, it will be configured with default behavior.
+   */
+  public fun compile(
+    @Language("kotlin") vararg sources: String,
+    block: Result.() -> Unit = {}
+  ): Result {
+    checkNotCompiled()
+    if (!anvilConfigured) {
+      // Configure with default behaviors
+      configureAnvil()
+    }
+    addSources(*sources)
+    isCompiled = true
+    return kotlinCompilation.compile()
+      .apply(block)
+      .apply { swapClassLoaderIfNecessary() }
+  }
+
+  // This workaround is needed until this issue is resolved:
+  // https://github.com/tschuchortdev/kotlin-compile-testing/issues/216
+  private fun Result.swapClassLoaderIfNecessary() {
+    if (!swapClassLoader) return
+
+    val newClassLoader = URLClassLoader(
+      kotlinCompilation.classpaths.plus(outputDirectory).map { it.toURI().toURL() }.toTypedArray(),
+      this::class.java.classLoader
+    )
+
+    Result::class.java.declaredFields
+      .single { it.name == "classLoader" }
+      .use { it.set(this, newClassLoader) }
+  }
+
+  companion object {
+    public operator fun invoke(): AnvilCompilation {
+      return AnvilCompilation(
+        KotlinCompilation().apply {
+          // Sensible default behaviors
+          inheritClassPath = true
+          jvmTarget = JvmTarget.JVM_1_8.description
+          verbose = false
+        }
+      )
+    }
+  }
+}
+
+/**
+ * Helpful for testing code generators in unit tests end to end.
+ *
+ * This covers common cases, but is built upon reusable logic in [AnvilCompilation] and
+ * [configureForAnvil]. Consider using those APIs if more advanced configuration is needed.
+ */
+@ExperimentalAnvilApi
+public fun compileAnvil(
+  @Language("kotlin") vararg sources: String,
+  enableDaggerAnnotationProcessor: Boolean = false,
+  generateDaggerFactories: Boolean = false,
+  generateDaggerFactoriesOnly: Boolean = false,
+  disableComponentMerging: Boolean = false,
+  allWarningsAsErrors: Boolean = true,
+  useIR: Boolean = true,
+  messageOutputStream: OutputStream = System.out,
+  workingDir: File? = null,
+  enableExperimentalAnvilApis: Boolean = true,
+  previousCompilationResult: Result? = null,
+  codeGenerators: List<CodeGenerator> = emptyList(),
+  block: Result.() -> Unit = { }
+): Result {
+  return AnvilCompilation()
+    .apply {
+      kotlinCompilation.apply {
+        this.allWarningsAsErrors = allWarningsAsErrors
+        this.messageOutputStream = messageOutputStream
+        if (workingDir != null) {
+          this.workingDir = workingDir
+        }
+      }
+
+      if (previousCompilationResult != null) {
+        addPreviousCompilationResult(previousCompilationResult)
+      }
+    }
+    .configureAnvil(
+      enableDaggerAnnotationProcessor = enableDaggerAnnotationProcessor,
+      generateDaggerFactories = generateDaggerFactories,
+      generateDaggerFactoriesOnly = generateDaggerFactoriesOnly,
+      disableComponentMerging = disableComponentMerging,
+      enableExperimentalAnvilApis = enableExperimentalAnvilApis,
+      codeGenerators = codeGenerators,
+    )
+    .useIR(useIR)
+    .compile(*sources)
+    .also(block)
+}
diff --git a/compiler-utils/src/testFixtures/java/com/squareup/anvil/compiler/internal/testing/TestUtils.kt b/compiler-utils/src/testFixtures/java/com/squareup/anvil/compiler/internal/testing/TestUtils.kt
index b5cd9988..81afc4c4 100644
--- a/compiler-utils/src/testFixtures/java/com/squareup/anvil/compiler/internal/testing/TestUtils.kt
+++ b/compiler-utils/src/testFixtures/java/com/squareup/anvil/compiler/internal/testing/TestUtils.kt
@@ -4,263 +4,12 @@ package com.squareup.anvil.compiler.internal.testing
 
 import com.google.common.truth.Truth.assertThat
 import com.squareup.anvil.annotations.ExperimentalAnvilApi
-import com.squareup.anvil.compiler.AnvilCommandLineProcessor
-import com.squareup.anvil.compiler.AnvilComponentRegistrar
 import com.squareup.anvil.compiler.internal.capitalize
-import com.tschuchort.compiletesting.KotlinCompilation
-import com.tschuchort.compiletesting.KotlinCompilation.Result
-import com.tschuchort.compiletesting.PluginOption
-import com.tschuchort.compiletesting.SourceFile
 import dagger.Component
 import dagger.Module
 import dagger.Subcomponent
-import dagger.internal.codegen.ComponentProcessor
-import org.intellij.lang.annotations.Language
-import org.jetbrains.kotlin.config.JvmTarget
-import java.io.File
-import java.io.OutputStream
-import java.net.URLClassLoader
-import java.nio.file.Files
 import kotlin.reflect.KClass
 
-/**
- * A simple API over a [KotlinCompilation] with extra configuration support for Anvil.
- */
-@ExperimentalAnvilApi
-public class AnvilCompilation internal constructor(
-  val kotlinCompilation: KotlinCompilation
-) {
-
-  private var isCompiled = false
-  private var anvilConfigured = false
-  private var swapClassLoader = false
-
-  /** Configures this the Anvil behavior of this compilation. */
-  @Suppress("SuspiciousCollectionReassignment")
-  @ExperimentalAnvilApi
-  public fun configureAnvil(
-    enableDaggerAnnotationProcessor: Boolean = false,
-    generateDaggerFactories: Boolean = false,
-    generateDaggerFactoriesOnly: Boolean = false,
-    disableComponentMerging: Boolean = false,
-    enableExperimentalAnvilApis: Boolean = true,
-  ) = apply {
-    checkNotCompiled()
-    anvilConfigured = true
-    kotlinCompilation.apply {
-      compilerPlugins = listOf(AnvilComponentRegistrar())
-      if (enableDaggerAnnotationProcessor) {
-        annotationProcessors = listOf(ComponentProcessor())
-      }
-
-      val anvilCommandLineProcessor = AnvilCommandLineProcessor()
-      commandLineProcessors = listOf(anvilCommandLineProcessor)
-
-      pluginOptions = listOf(
-        PluginOption(
-          pluginId = anvilCommandLineProcessor.pluginId,
-          optionName = "src-gen-dir",
-          optionValue = File(workingDir, "build/anvil").absolutePath
-        ),
-        PluginOption(
-          pluginId = anvilCommandLineProcessor.pluginId,
-          optionName = "generate-dagger-factories",
-          optionValue = generateDaggerFactories.toString()
-        ),
-        PluginOption(
-          pluginId = anvilCommandLineProcessor.pluginId,
-          optionName = "generate-dagger-factories-only",
-          optionValue = generateDaggerFactoriesOnly.toString()
-        ),
-        PluginOption(
-          pluginId = anvilCommandLineProcessor.pluginId,
-          optionName = "disable-component-merging",
-          optionValue = disableComponentMerging.toString()
-        )
-      )
-
-      if (enableExperimentalAnvilApis) {
-        kotlincArguments += listOf(
-          "-Xopt-in=kotlin.RequiresOptIn",
-          "-Xopt-in=com.squareup.anvil.annotations.ExperimentalAnvilApi"
-        )
-      }
-    }
-  }
-
-  /**
-   * Helpful shim to configure both [KotlinCompilation.useIR] and [KotlinCompilation.useOldBackend]
-   * accordingly.
-   */
-  public fun useIR(useIR: Boolean) = apply {
-    checkNotCompiled()
-    kotlinCompilation.useIR = useIR
-    kotlinCompilation.useOldBackend = !useIR
-
-    if (!useIR) {
-      // To suppress:
-      //  w: Language version is automatically inferred to 1.5 when using the old JVM backend.
-      //    Consider specifying -language-version explicitly, or remove -Xuse-old-backend
-      //  e: warnings found and -Werror specified
-      kotlinCompilation.languageVersion = "1.5"
-
-      // To suppress:
-      //  w: -Xuse-old-backend is deprecated and will be removed in a future release
-      //  e: warnings found and -Werror specified
-      kotlinCompilation.allWarningsAsErrors = false
-    }
-  }
-
-  /** Adds the given sources to this compilation with their packages and names inferred. */
-  public fun addSources(@Language("kotlin") vararg sources: String) = apply {
-    checkNotCompiled()
-    kotlinCompilation.sources += sources.mapIndexed { index, content ->
-      val packageDir = content.lines()
-        .firstOrNull { it.trim().startsWith("package ") }
-        ?.substringAfter("package ")
-        ?.replace('.', '/')
-        ?.let { "$it/" }
-        ?: ""
-
-      val name = "${kotlinCompilation.workingDir.absolutePath}/sources/src/main/java/" +
-        "$packageDir/Source$index.kt"
-
-      Files.createDirectories(File(name).parentFile.toPath())
-
-      SourceFile.kotlin(name, contents = content, trimIndent = true)
-    }
-  }
-
-  public fun addPreviousCompilationResult(result: Result) = apply {
-    checkNotCompiled()
-    kotlinCompilation.classpaths += result.outputDirectory
-    swapClassLoader = true
-  }
-
-  /**
-   * Returns an Anvil-generated file with the given [packageName] and [fileName] from its expected
-   * path.
-   */
-  public fun generatedAnvilFile(
-    packageName: String,
-    fileName: String
-  ): File {
-    check(isCompiled) {
-      "No compilation run yet! Call compile() first."
-    }
-    return File(
-      kotlinCompilation.workingDir,
-      "build/anvil/${packageName.replace('.', File.separatorChar)}/$fileName.kt"
-    )
-      .apply {
-        check(exists()) {
-          "Generated file not found!"
-        }
-      }
-  }
-
-  private fun checkNotCompiled() {
-    check(!isCompiled) {
-      "Already compiled! Create a new compilation if you want to compile again."
-    }
-  }
-
-  /**
-   * Compiles the underlying [KotlinCompilation]. Note that if [configureAnvil] has not been called
-   * prior to this, it will be configured with default behavior.
-   */
-  public fun compile(
-    @Language("kotlin") vararg sources: String,
-    block: Result.() -> Unit = {}
-  ): Result {
-    checkNotCompiled()
-    if (!anvilConfigured) {
-      // Configure with default behaviors
-      configureAnvil()
-    }
-    addSources(*sources)
-    isCompiled = true
-    return kotlinCompilation.compile()
-      .apply(block)
-      .apply { swapClassLoaderIfNecessary() }
-  }
-
-  // This workaround is needed until this issue is resolved:
-  // https://github.com/tschuchortdev/kotlin-compile-testing/issues/216
-  private fun Result.swapClassLoaderIfNecessary() {
-    if (!swapClassLoader) return
-
-    val newClassLoader = URLClassLoader(
-      kotlinCompilation.classpaths.plus(outputDirectory).map { it.toURI().toURL() }.toTypedArray(),
-      this::class.java.classLoader
-    )
-
-    Result::class.java.declaredFields
-      .single { it.name == "classLoader" }
-      .use { it.set(this, newClassLoader) }
-  }
-
-  companion object {
-    public operator fun invoke(): AnvilCompilation {
-      return AnvilCompilation(
-        KotlinCompilation().apply {
-          // Sensible default behaviors
-          inheritClassPath = true
-          jvmTarget = JvmTarget.JVM_1_8.description
-          verbose = false
-        }
-      )
-    }
-  }
-}
-
-/**
- * Helpful for testing code generators in unit tests end to end.
- *
- * This covers common cases, but is built upon reusable logic in [AnvilCompilation] and
- * [configureForAnvil]. Consider using those APIs if more advanced configuration is needed.
- */
-@ExperimentalAnvilApi
-public fun compileAnvil(
-  @Language("kotlin") vararg sources: String,
-  enableDaggerAnnotationProcessor: Boolean = false,
-  generateDaggerFactories: Boolean = false,
-  generateDaggerFactoriesOnly: Boolean = false,
-  disableComponentMerging: Boolean = false,
-  allWarningsAsErrors: Boolean = true,
-  useIR: Boolean = true,
-  messageOutputStream: OutputStream = System.out,
-  workingDir: File? = null,
-  enableExperimentalAnvilApis: Boolean = true,
-  previousCompilationResult: Result? = null,
-  block: Result.() -> Unit = { }
-): Result {
-  return AnvilCompilation()
-    .apply {
-      kotlinCompilation.apply {
-        this.allWarningsAsErrors = allWarningsAsErrors
-        this.messageOutputStream = messageOutputStream
-        if (workingDir != null) {
-          this.workingDir = workingDir
-        }
-      }
-
-      if (previousCompilationResult != null) {
-        addPreviousCompilationResult(previousCompilationResult)
-      }
-    }
-    .configureAnvil(
-      enableDaggerAnnotationProcessor = enableDaggerAnnotationProcessor,
-      generateDaggerFactories = generateDaggerFactories,
-      generateDaggerFactoriesOnly = generateDaggerFactoriesOnly,
-      disableComponentMerging = disableComponentMerging,
-      enableExperimentalAnvilApis = enableExperimentalAnvilApis,
-    )
-    .useIR(useIR)
-    .compile(*sources)
-    .also(block)
-}
-
 @ExperimentalAnvilApi
 public fun Class<*>.moduleFactoryClass(
   providerMethodName: String,
diff --git a/compiler/src/test/java/com/squareup/anvil/compiler/TestUtils.kt b/compiler/src/test/java/com/squareup/anvil/compiler/TestUtils.kt
index ac8e0b5a..13addbbb 100644
--- a/compiler/src/test/java/com/squareup/anvil/compiler/TestUtils.kt
+++ b/compiler/src/test/java/com/squareup/anvil/compiler/TestUtils.kt
@@ -3,6 +3,7 @@ package com.squareup.anvil.compiler
 import com.google.common.truth.ComparableSubject
 import com.google.common.truth.Truth.assertThat
 import com.squareup.anvil.annotations.MergeComponent
+import com.squareup.anvil.compiler.api.CodeGenerator
 import com.squareup.anvil.compiler.internal.capitalize
 import com.squareup.anvil.compiler.internal.testing.compileAnvil
 import com.squareup.anvil.compiler.internal.testing.generatedClassesString
@@ -19,6 +20,7 @@ internal fun compile(
   @Language("kotlin") vararg sources: String,
   previousCompilationResult: Result? = null,
   enableDaggerAnnotationProcessor: Boolean = false,
+  codeGenerators: List<CodeGenerator> = emptyList(),
   block: Result.() -> Unit = { }
 ): Result = compileAnvil(
   sources = sources,
@@ -26,6 +28,7 @@ internal fun compile(
   allWarningsAsErrors = WARNINGS_AS_ERRORS,
   previousCompilationResult = previousCompilationResult,
   enableDaggerAnnotationProcessor = enableDaggerAnnotationProcessor,
+  codeGenerators = codeGenerators,
   block = block
 )