square · matejdro · Jul 21, 2023 · Jul 21, 2023 · Jul 21, 2023 · Jul 21, 2023
diff --git a/README.md b/README.md
@@ -237,6 +237,17 @@ sources either.
 When you enable this feature, don't forget to remove the Dagger annotation processor. You should
 keep all other dependencies.
 
+If you want to only enable Dagger factory generation on some source sets 
+(such as only in main source set or only in tests), you can add those source sets
+to the whitelist:
+
+```groovy
+anvil {
+  generateDaggerFactoriesSourceSetWhitelist.add("main")
+}
+```
+
+
 ## Extending Anvil
 
 Every codebase has its own dependency injection patterns where certain code structures need to be

diff --git a/compiler-api/src/main/java/com/squareup/anvil/compiler/api/AnvilContext.kt b/compiler-api/src/main/java/com/squareup/anvil/compiler/api/AnvilContext.kt
@@ -34,6 +34,13 @@ public interface AnvilContext {
    */
   public val generateFactoriesOnly: Boolean
 
+  /**
+   * List of absolute paths that anvil should generate Factory classes in.
+   * When empty or ommited, Anvil applies no filtering and
+   * generates factories in all available files.
+   */
+  public val generateFactoriesPathWhitelist: List<String>
+
   /**
    * Enabling this indicates that only code generation should run and no component merging should
    * run. This is useful for cases where you want to use `@ContributesTo`, `@ContributesBinding`,

diff --git a/...ls/src/testFixtures/java/com/squareup/anvil/compiler/internal/testing/AnvilCompilation.kt b/...ls/src/testFixtures/java/com/squareup/anvil/compiler/internal/testing/AnvilCompilation.kt
@@ -5,6 +5,9 @@ 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.squareup.anvil.compiler.internal.testing.TestWhitelistType.BLACKLISTED
+import com.squareup.anvil.compiler.internal.testing.TestWhitelistType.EMPTY
+import com.squareup.anvil.compiler.internal.testing.TestWhitelistType.WHITELISTED
 import com.tschuchort.compiletesting.KotlinCompilation
 import com.tschuchort.compiletesting.KotlinCompilation.Result
 import com.tschuchort.compiletesting.PluginOption
@@ -35,6 +38,7 @@ public class AnvilCompilation internal constructor(
     enableDaggerAnnotationProcessor: Boolean = false,
     generateDaggerFactories: Boolean = false,
     generateDaggerFactoriesOnly: Boolean = false,
+    generateDaggerFactoriesPathWhitelist: List<String> = emptyList(),
     disableComponentMerging: Boolean = false,
     enableExperimentalAnvilApis: Boolean = true,
     codeGenerators: List<CodeGenerator> = emptyList(),
@@ -71,6 +75,11 @@ public class AnvilCompilation internal constructor(
           optionName = "generate-dagger-factories",
           optionValue = generateDaggerFactories.toString()
         ),
+        PluginOption(
+          pluginId = anvilCommandLineProcessor.pluginId,
+          optionName = "generate-dagger-factories-path-whitelist",
+          optionValue = generateDaggerFactoriesPathWhitelist.joinToString()
+        ),
         PluginOption(
           pluginId = anvilCommandLineProcessor.pluginId,
           optionName = "generate-dagger-factories-only",
@@ -195,6 +204,7 @@ public fun compileAnvil(
   @Language("kotlin") vararg sources: String,
   enableDaggerAnnotationProcessor: Boolean = false,
   generateDaggerFactories: Boolean = false,
+  generateDaggerFactoriesWhitelist: TestWhitelistType = EMPTY,
   generateDaggerFactoriesOnly: Boolean = false,
   disableComponentMerging: Boolean = false,
   allWarningsAsErrors: Boolean = true,
@@ -207,6 +217,8 @@ public fun compileAnvil(
   jvmTarget: JvmTarget? = null,
   block: Result.() -> Unit = { },
 ): Result {
+  val whitelist: List<String>
+
   return AnvilCompilation()
     .apply {
       kotlinCompilation.apply {
@@ -227,10 +239,17 @@ public fun compileAnvil(
       if (previousCompilationResult != null) {
         addPreviousCompilationResult(previousCompilationResult)
       }
+
+      whitelist = when (generateDaggerFactoriesWhitelist) {
+        EMPTY -> emptyList()
+        WHITELISTED -> listOf(kotlinCompilation.workingDir.absolutePath)
+        BLACKLISTED -> listOf("some-other-folder")
+      }
     }
     .configureAnvil(
       enableDaggerAnnotationProcessor = enableDaggerAnnotationProcessor,
       generateDaggerFactories = generateDaggerFactories,
+      generateDaggerFactoriesPathWhitelist = whitelist,
       generateDaggerFactoriesOnly = generateDaggerFactoriesOnly,
       disableComponentMerging = disableComponentMerging,
       enableExperimentalAnvilApis = enableExperimentalAnvilApis,
@@ -239,3 +258,9 @@ public fun compileAnvil(
     .compile(*sources)
     .also(block)
 }
+
+enum class TestWhitelistType {
+  EMPTY,
+  WHITELISTED,
+  BLACKLISTED
+}
diff --git a/compiler/src/main/java/com/squareup/anvil/compiler/AnvilCommandLineProcessor.kt b/compiler/src/main/java/com/squareup/anvil/compiler/AnvilCommandLineProcessor.kt
@@ -18,6 +18,10 @@ internal const val generateDaggerFactoriesOnlyName = "generate-dagger-factories-
 internal val generateDaggerFactoriesOnlyKey =
   CompilerConfigurationKey.create<Boolean>("anvil $generateDaggerFactoriesOnlyName")
 
+internal const val generateDaggerFactoriesPathWhitelistName = "generate-dagger-factories-path-whitelist"
+internal val generateDaggerFactoriesPathWhitelistKey =
+  CompilerConfigurationKey.create<List<String>>("anvil $generateDaggerFactoriesPathWhitelistName")
+
 internal const val disableComponentMergingName = "disable-component-merging"
 internal val disableComponentMergingKey =
   CompilerConfigurationKey.create<Boolean>("anvil $disableComponentMergingName")
@@ -54,6 +58,15 @@ class AnvilCommandLineProcessor : CommandLineProcessor {
       required = false,
       allowMultipleOccurrences = false
     ),
+    CliOption(
+      optionName = generateDaggerFactoriesPathWhitelistName,
+      valueDescription = "Colon separated list",
+      description = "List of absolute paths that anvil should generate Factory classes in. " +
+      "When empty or ommited, Anvil applies no filtering and generates factories " +
+        "in all available files",
+      required = false,
+      allowMultipleOccurrences = false
+    ),
     CliOption(
       optionName = disableComponentMergingName,
       valueDescription = "<true|false>",
@@ -77,6 +90,11 @@ class AnvilCommandLineProcessor : CommandLineProcessor {
         configuration.put(generateDaggerFactoriesOnlyKey, value.toBoolean())
       disableComponentMergingName ->
         configuration.put(disableComponentMergingKey, value.toBoolean())
+      generateDaggerFactoriesPathWhitelistName ->
+        configuration.put(
+          generateDaggerFactoriesPathWhitelistKey,
+          value.split(":").filter { it.isNotBlank() }
+        )
     }
   }
 }
diff --git a/compiler/src/main/java/com/squareup/anvil/compiler/CommandLineOptions.kt b/compiler/src/main/java/com/squareup/anvil/compiler/CommandLineOptions.kt
@@ -5,13 +5,15 @@ import org.jetbrains.kotlin.config.CompilerConfiguration
 class CommandLineOptions private constructor(
   val generateFactories: Boolean,
   val generateFactoriesOnly: Boolean,
+  val generateFactoriesPathWhitelist: List<String>,
   val disableComponentMerging: Boolean,
 ) {
   companion object {
     val CompilerConfiguration.commandLineOptions: CommandLineOptions
       get() = CommandLineOptions(
         generateFactories = get(generateDaggerFactoriesKey, false),
         generateFactoriesOnly = get(generateDaggerFactoriesOnlyKey, false),
+        generateFactoriesPathWhitelist = get(generateDaggerFactoriesPathWhitelistKey, emptyList()),
         disableComponentMerging = get(disableComponentMergingKey, false)
       )
   }

diff --git a/compiler/src/main/java/com/squareup/anvil/compiler/codegen/CodeGenerationExtension.kt b/compiler/src/main/java/com/squareup/anvil/compiler/codegen/CodeGenerationExtension.kt
@@ -8,6 +8,7 @@ import com.squareup.anvil.compiler.codegen.reference.RealAnvilModuleDescriptor
 import org.jetbrains.kotlin.analyzer.AnalysisResult
 import org.jetbrains.kotlin.analyzer.AnalysisResult.RetryWithAdditionalRoots
 import org.jetbrains.kotlin.com.intellij.openapi.project.Project
+import org.jetbrains.kotlin.com.intellij.openapi.util.Key
 import org.jetbrains.kotlin.com.intellij.psi.PsiManager
 import org.jetbrains.kotlin.com.intellij.testFramework.LightVirtualFile
 import org.jetbrains.kotlin.container.ComponentProvider
@@ -93,7 +94,9 @@ internal class CodeGenerationExtension(
             content
           )
 
-          psiManager.findFile(virtualFile)
+          psiManager.findFile(virtualFile)?.also {
+            it.putUserData(IS_GENERATED_BY_ANVIL, true)
+          }
         }
         .filterIsInstance<KtFile>()
         .also { anvilModule.addFiles(it) }
@@ -160,3 +163,5 @@ internal class CodeGenerationExtension(
     )
   }
 }
+
+val IS_GENERATED_BY_ANVIL = Key<Boolean>("generated_by_anvil")
diff --git a/compiler/src/main/java/com/squareup/anvil/compiler/codegen/RealAnvilContext.kt b/compiler/src/main/java/com/squareup/anvil/compiler/codegen/RealAnvilContext.kt
@@ -7,6 +7,7 @@ import org.jetbrains.kotlin.descriptors.ModuleDescriptor
 internal data class RealAnvilContext(
   override val generateFactories: Boolean,
   override val generateFactoriesOnly: Boolean,
+  override val generateFactoriesPathWhitelist: List<String>,
   override val disableComponentMerging: Boolean,
   override val module: ModuleDescriptor
 ) : AnvilContext
@@ -16,6 +17,7 @@ internal fun CommandLineOptions.toAnvilContext(
 ): RealAnvilContext = RealAnvilContext(
   generateFactories = generateFactories,
   generateFactoriesOnly = generateFactoriesOnly,
+  generateFactoriesPathWhitelist = generateFactoriesPathWhitelist,
   disableComponentMerging = disableComponentMerging,
   module = module
 )
diff --git a/...iler/src/main/java/com/squareup/anvil/compiler/codegen/dagger/AssistedFactoryGenerator.kt b/...iler/src/main/java/com/squareup/anvil/compiler/codegen/dagger/AssistedFactoryGenerator.kt
@@ -43,11 +43,9 @@ import java.io.File
 import javax.inject.Provider
 
 @AutoService(CodeGenerator::class)
-internal class AssistedFactoryGenerator : PrivateCodeGenerator() {
+internal class AssistedFactoryGenerator : DaggerFactoryGenerator() {
 
-  override fun isApplicable(context: AnvilContext) = context.generateFactories
-
-  override fun generateCodePrivate(
+  override fun generateCodeInDaggerFactoryWhitelistedFiles(
     codeGenDir: File,
     module: ModuleDescriptor,
     projectFiles: Collection<KtFile>

diff --git a/compiler/src/main/java/com/squareup/anvil/compiler/codegen/dagger/AssistedInjectGenerator.kt b/compiler/src/main/java/com/squareup/anvil/compiler/codegen/dagger/AssistedInjectGenerator.kt
@@ -38,11 +38,9 @@ import java.io.File
  * ```
  */
 @AutoService(CodeGenerator::class)
-internal class AssistedInjectGenerator : PrivateCodeGenerator() {
+internal class AssistedInjectGenerator : DaggerFactoryGenerator() {
 
-  override fun isApplicable(context: AnvilContext) = context.generateFactories
-
-  override fun generateCodePrivate(
+  override fun generateCodeInDaggerFactoryWhitelistedFiles(
     codeGenDir: File,
     module: ModuleDescriptor,
     projectFiles: Collection<KtFile>

diff --git a/compiler/src/main/java/com/squareup/anvil/compiler/codegen/dagger/BindsMethodValidator.kt b/compiler/src/main/java/com/squareup/anvil/compiler/codegen/dagger/BindsMethodValidator.kt
@@ -25,11 +25,9 @@ import java.io.File
  * when Anvil generates Dagger factories.
  */
 @AutoService(CodeGenerator::class)
-internal class BindsMethodValidator : PrivateCodeGenerator() {
+internal class BindsMethodValidator : DaggerFactoryGenerator() {
 
-  override fun isApplicable(context: AnvilContext) = context.generateFactories
-
-  override fun generateCodePrivate(
+  override fun generateCodeInDaggerFactoryWhitelistedFiles(
     codeGenDir: File,
     module: ModuleDescriptor,
     projectFiles: Collection<KtFile>

diff --git a/compiler/src/main/java/com/squareup/anvil/compiler/codegen/dagger/ComponentDetectorCheck.kt b/compiler/src/main/java/com/squareup/anvil/compiler/codegen/dagger/ComponentDetectorCheck.kt
@@ -12,11 +12,8 @@ import org.jetbrains.kotlin.psi.KtFile
 import java.io.File
 
 @AutoService(CodeGenerator::class)
-internal class ComponentDetectorCheck : PrivateCodeGenerator() {
-
-  override fun isApplicable(context: AnvilContext) = context.generateFactories
-
-  override fun generateCodePrivate(
+internal class ComponentDetectorCheck : DaggerFactoryGenerator() {
+  override fun generateCodeInDaggerFactoryWhitelistedFiles(
     codeGenDir: File,
     module: ModuleDescriptor,
     projectFiles: Collection<KtFile>

diff --git a/compiler/src/main/java/com/squareup/anvil/compiler/codegen/dagger/DaggerFactoryGenerator.kt b/compiler/src/main/java/com/squareup/anvil/compiler/codegen/dagger/DaggerFactoryGenerator.kt
@@ -0,0 +1,53 @@
+package com.squareup.anvil.compiler.codegen.dagger
+
+import com.squareup.anvil.compiler.api.AnvilContext
+import com.squareup.anvil.compiler.codegen.IS_GENERATED_BY_ANVIL
+import com.squareup.anvil.compiler.codegen.PrivateCodeGenerator
+import org.jetbrains.kotlin.com.intellij.openapi.vfs.VirtualFile
+import org.jetbrains.kotlin.com.intellij.openapi.vfs.local.CoreLocalVirtualFile
+import org.jetbrains.kotlin.com.intellij.testFramework.LightVirtualFile
+import org.jetbrains.kotlin.descriptors.ModuleDescriptor
+import org.jetbrains.kotlin.psi.KtFile
+import java.io.File
+
+internal abstract class DaggerFactoryGenerator: PrivateCodeGenerator() {
+  protected var pathWhitelist = emptyList<String>()
+
+  override fun isApplicable(context: AnvilContext): Boolean {
+    pathWhitelist = context.generateFactoriesPathWhitelist
+
+    return context.generateFactories
+  }
+
+  abstract fun generateCodeInDaggerFactoryWhitelistedFiles(
+    codeGenDir: File,
+    module: ModuleDescriptor,
+    projectFiles: Collection<KtFile>
+  )
+
+  override fun generateCodePrivate(
+    codeGenDir: File,
+    module: ModuleDescriptor,
+    projectFiles: Collection<KtFile>
+  ) {
+      return generateCodeInDaggerFactoryWhitelistedFiles(
+        codeGenDir,
+        module,
+        projectFiles.filter { ktFile ->
+          ktFile.getUserData(IS_GENERATED_BY_ANVIL) == true ||
+            pathWhitelist.isEmpty() ||
+            pathWhitelist.any { ktFile.virtualFilePath.startsWith(it) }
+        }
+      )
+  }
+
+  private fun VirtualFile.absolutePath(): String {
+      return if (this is LightVirtualFile) {
+        originalFile.path
+      } else if (this is CoreLocalVirtualFile) {
+        this.path
+      } else {
+        throw UnsupportedOperationException("Cannot determine path of $this (${this.javaClass})")
+      }
+  }
+}
diff --git a/...main/java/com/squareup/anvil/compiler/codegen/dagger/InjectConstructorFactoryGenerator.kt b/...main/java/com/squareup/anvil/compiler/codegen/dagger/InjectConstructorFactoryGenerator.kt
@@ -31,11 +31,9 @@ import org.jetbrains.kotlin.psi.KtFile
 import java.io.File
 
 @AutoService(CodeGenerator::class)
-internal class InjectConstructorFactoryGenerator : PrivateCodeGenerator() {
+internal class InjectConstructorFactoryGenerator : DaggerFactoryGenerator() {
 
-  override fun isApplicable(context: AnvilContext) = context.generateFactories
-
-  override fun generateCodePrivate(
+  override fun generateCodeInDaggerFactoryWhitelistedFiles(
     codeGenDir: File,
     module: ModuleDescriptor,
     projectFiles: Collection<KtFile>

diff --git a/compiler/src/main/java/com/squareup/anvil/compiler/codegen/dagger/MapKeyCreatorGenerator.kt b/compiler/src/main/java/com/squareup/anvil/compiler/codegen/dagger/MapKeyCreatorGenerator.kt
@@ -56,11 +56,9 @@ import kotlin.reflect.KClass
  * Implemented from eyeballing https://github.com/google/dagger/blob/b5990a0641a7860b760aa9055b90a99d06186af6/javatests/dagger/internal/codegen/MapKeyProcessorTest.java
  */
 @AutoService(CodeGenerator::class)
-internal class MapKeyCreatorGenerator : PrivateCodeGenerator() {
+internal class MapKeyCreatorGenerator : DaggerFactoryGenerator() {
 
-  override fun isApplicable(context: AnvilContext) = context.generateFactories
-
-  override fun generateCodePrivate(
+  override fun generateCodeInDaggerFactoryWhitelistedFiles(
     codeGenDir: File,
     module: ModuleDescriptor,
     projectFiles: Collection<KtFile>

diff --git a/...iler/src/main/java/com/squareup/anvil/compiler/codegen/dagger/MembersInjectorGenerator.kt b/...iler/src/main/java/com/squareup/anvil/compiler/codegen/dagger/MembersInjectorGenerator.kt
@@ -33,11 +33,9 @@ import org.jetbrains.kotlin.psi.KtFile
 import java.io.File
 
 @AutoService(CodeGenerator::class)
-internal class MembersInjectorGenerator : PrivateCodeGenerator() {
+internal class MembersInjectorGenerator : DaggerFactoryGenerator() {
 
-  override fun isApplicable(context: AnvilContext) = context.generateFactories
-
-  override fun generateCodePrivate(
+  override fun generateCodeInDaggerFactoryWhitelistedFiles(
     codeGenDir: File,
     module: ModuleDescriptor,
     projectFiles: Collection<KtFile>

diff --git a/...rc/main/java/com/squareup/anvil/compiler/codegen/dagger/ProvidesMethodFactoryGenerator.kt b/...rc/main/java/com/squareup/anvil/compiler/codegen/dagger/ProvidesMethodFactoryGenerator.kt
@@ -43,11 +43,9 @@ import org.jetbrains.kotlin.psi.KtFile
 import java.io.File
 
 @AutoService(CodeGenerator::class)
-internal class ProvidesMethodFactoryGenerator : PrivateCodeGenerator() {
+internal class ProvidesMethodFactoryGenerator : DaggerFactoryGenerator() {
 
-  override fun isApplicable(context: AnvilContext) = context.generateFactories
-
-  override fun generateCodePrivate(
+  override fun generateCodeInDaggerFactoryWhitelistedFiles(
     codeGenDir: File,
     module: ModuleDescriptor,
     projectFiles: Collection<KtFile>