Repository: square/AssistedInject (kotlin)

Issue:
The generated module should be annotated with `@DisableInstallInCheck`
Now that Hilt is live it would be great if the generated module could be annotated with `@DisableInstallInCheck`

Pull request:
This fixes #150

Patch:
diff --git a/assisted-inject-processor-dagger2/src/main/java/com/squareup/inject/assisted/dagger2/processor/AssistedInjectDagger2Processor.kt b/assisted-inject-processor-dagger2/src/main/java/com/squareup/inject/assisted/dagger2/processor/AssistedInjectDagger2Processor.kt
index 7438e37..07871ce 100644
--- a/assisted-inject-processor-dagger2/src/main/java/com/squareup/inject/assisted/dagger2/processor/AssistedInjectDagger2Processor.kt
+++ b/assisted-inject-processor-dagger2/src/main/java/com/squareup/inject/assisted/dagger2/processor/AssistedInjectDagger2Processor.kt
@@ -18,6 +18,7 @@ package com.squareup.inject.assisted.dagger2.processor
 import com.google.auto.service.AutoService
 import com.squareup.inject.assisted.AssistedInject
 import com.squareup.inject.assisted.dagger2.AssistedModule
+import com.squareup.inject.assisted.dagger2.processor.internal.generateHiltAnnotation
 import com.squareup.inject.assisted.processor.internal.MirrorValue
 import com.squareup.inject.assisted.processor.internal.applyEach
 import com.squareup.inject.assisted.processor.internal.cast
@@ -117,9 +118,16 @@ class AssistedInjectDagger2Processor : AbstractProcessor() {
         val referencesGeneratedModule = includes
             .map { it.toTypeName() }
             .any { it == generatedModuleName }
-        if (!referencesGeneratedModule) {
-          error("@AssistedModule's @Module must include ${generatedModuleName.simpleName()}",
+        if (userModule.hasAnnotation("dagger.hilt.InstallIn")) {
+          if (referencesGeneratedModule) {
+            error("@AssistedModule's @Module must not include ${generatedModuleName.simpleName()} if @InstallIn is used",
               userModule)
+          }
+        } else {
+          if (!referencesGeneratedModule) {
+            error("@AssistedModule's @Module must include ${generatedModuleName.simpleName()}",
+              userModule)
+          }
         }
       }
     }
@@ -160,9 +168,11 @@ class AssistedInjectDagger2Processor : AbstractProcessor() {
         .map { (target, factory) -> target.asType().toTypeName() to factory.toClassName() }
         .toMap(TreeMap())
     val public = Modifier.PUBLIC in moduleType.modifiers
-    val generatedAnnotation = createGeneratedAnnotation(sourceVersion, elements)
-    return AssistedInjectionModule(moduleName, public, targetNameToFactoryNames,
-        generatedAnnotation)
+    val extraAnnotations = listOfNotNull(
+      createGeneratedAnnotation(sourceVersion, elements),
+      generateHiltAnnotation(moduleType, elements)
+    )
+    return AssistedInjectionModule(moduleName, public, targetNameToFactoryNames, extraAnnotations)
   }
 
   private fun writeAssistedModule(
diff --git a/assisted-inject-processor-dagger2/src/main/java/com/squareup/inject/assisted/dagger2/processor/AssistedInjectionModule.kt b/assisted-inject-processor-dagger2/src/main/java/com/squareup/inject/assisted/dagger2/processor/AssistedInjectionModule.kt
index 670c266..05f6bc8 100644
--- a/assisted-inject-processor-dagger2/src/main/java/com/squareup/inject/assisted/dagger2/processor/AssistedInjectionModule.kt
+++ b/assisted-inject-processor-dagger2/src/main/java/com/squareup/inject/assisted/dagger2/processor/AssistedInjectionModule.kt
@@ -18,8 +18,8 @@ data class AssistedInjectionModule(
   val moduleName: ClassName,
   val public: Boolean,
   val targetNameToFactoryName: Map<TypeName, ClassName>,
-  /** An optional `@Generated` annotation marker. */
-  val generatedAnnotation: AnnotationSpec? = null
+  /** An optional list of extra annotations markers, e.g `@Generated`. */
+  val extraAnnotations: List<AnnotationSpec> = emptyList()
 ) {
   val generatedType = moduleName.assistedInjectModuleName()
 
@@ -27,8 +27,8 @@ data class AssistedInjectionModule(
     return TypeSpec.classBuilder(generatedType)
         .addAnnotation(MODULE)
         .apply {
-          if (generatedAnnotation != null) {
-            addAnnotation(generatedAnnotation)
+          for (annotation in extraAnnotations) {
+            addAnnotation(annotation)
           }
         }
         .addModifiers(ABSTRACT)
diff --git a/assisted-inject-processor-dagger2/src/main/java/com/squareup/inject/assisted/dagger2/processor/internal/hilt.kt b/assisted-inject-processor-dagger2/src/main/java/com/squareup/inject/assisted/dagger2/processor/internal/hilt.kt
new file mode 100644
index 0000000..cc7b2a7
--- /dev/null
+++ b/assisted-inject-processor-dagger2/src/main/java/com/squareup/inject/assisted/dagger2/processor/internal/hilt.kt
@@ -0,0 +1,37 @@
+/*
+ * Copyright (C) 2020 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
+ *
+ *      http://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.inject.assisted.dagger2.processor.internal
+
+import com.squareup.inject.assisted.processor.internal.getAnnotation
+import com.squareup.javapoet.AnnotationSpec
+import com.squareup.javapoet.ClassName
+import javax.lang.model.element.TypeElement
+import javax.lang.model.util.Elements
+
+internal fun generateHiltAnnotation(moduleType: TypeElement, elements: Elements): AnnotationSpec? {
+  val installIn = moduleType.getAnnotation("dagger.hilt.InstallIn")
+  if (installIn != null) {
+    // If the @AssistedModule class also has an @InstallIn we just mirror the annotation.
+    return AnnotationSpec.get(installIn)
+  }
+  // If there is no @InstallIn annotation we add the @DisableInstallInCheck annotation, if
+  // it exists.
+  val disableInstallInCheck = ClassName.get("dagger.hilt.migration", "DisableInstallInCheck")
+  if (elements.getTypeElement(disableInstallInCheck.reflectionName()) != null) {
+    return AnnotationSpec.builder(disableInstallInCheck).build()
+  }
+  return null
+}


Test patch:
diff --git a/assisted-inject-processor-dagger2/src/test/java/com/squareup/inject/assisted/dagger2/processor/AssistedInjectDagger2ProcessorTest.kt b/assisted-inject-processor-dagger2/src/test/java/com/squareup/inject/assisted/dagger2/processor/AssistedInjectDagger2ProcessorTest.kt
index 676c973..4f4a68d 100644
--- a/assisted-inject-processor-dagger2/src/test/java/com/squareup/inject/assisted/dagger2/processor/AssistedInjectDagger2ProcessorTest.kt
+++ b/assisted-inject-processor-dagger2/src/test/java/com/squareup/inject/assisted/dagger2/processor/AssistedInjectDagger2ProcessorTest.kt
@@ -350,6 +350,67 @@ class AssistedInjectDagger2ProcessorTest {
         .generatesSources(expected)
   }
 
+  @Test fun installInAnnotationCopiedToGeneratedModule() {
+    val test = JavaFileObjects.forSourceString("test.Test", """
+      package test;
+
+      import com.squareup.inject.assisted.Assisted;
+      import com.squareup.inject.assisted.AssistedInject;
+
+      class Test {
+        @AssistedInject
+        Test(Long foo, @Assisted String bar) {}
+
+        @AssistedInject.Factory
+        interface Factory {}
+      }
+    """)
+    val testFactory = JavaFileObjects.forSourceString("test.Test_AssistedFactory", """
+      package test;
+
+      class Test_AssistedFactory implements Test.Factory {}
+    """)
+    val module = JavaFileObjects.forSourceString("test.TestModule", """
+      package test;
+
+      import com.squareup.inject.assisted.dagger2.AssistedModule;
+      import dagger.Module;
+      import dagger.hilt.InstallIn;
+
+      @Module
+      @AssistedModule
+      @InstallIn(String.class)
+      abstract class TestModule {}
+    """)
+
+    val expected = JavaFileObjects.forSourceString("test.AssistedInject_TestModule", """
+      package test;
+
+      import dagger.Binds;
+      import dagger.Module;
+      import dagger.hilt.InstallIn;
+      import java.lang.String;
+      import $GENERATED_TYPE;
+
+      @Module
+      $GENERATED_ANNOTATION
+      @InstallIn(String.class)
+      abstract class AssistedInject_TestModule {
+        private AssistedInject_TestModule() {}
+
+        @Binds abstract Test.Factory bind_test_Test(Test_AssistedFactory factory);
+      }
+    """)
+
+    assertAbout(javaSources())
+        .that(listOf(hiltInstallIn, test, testFactory, module))
+        .processedWith(AssistedInjectDagger2Processor())
+        .compilesWithoutError()
+        .and()
+        .generatesSources(expected)
+  }
+
+
   @Test fun moduleWithoutModuleAnnotationFails() {
     val moduleOne = JavaFileObjects.forSourceString("test.OneModule", """
       package test;
@@ -464,6 +525,29 @@ class AssistedInjectDagger2ProcessorTest {
         .`in`(moduleOne).onLine(9)
   }
 
+  @Test fun moduleWithInstallInWithIncludeFails() {
+    val module = JavaFileObjects.forSourceString("test.TestModule", """
+      package test;
+
+      import com.squareup.inject.assisted.dagger2.AssistedModule;
+      import dagger.Module;
+      import dagger.hilt.InstallIn;
+
+      @AssistedModule
+      @InstallIn(String.class)
+      @Module(includes = AssistedInject_TestModule.class)
+      abstract class TestModule {}
+    """)
+
+    assertAbout(javaSources())
+        .that(listOf(hiltInstallIn, module))
+        .processedWith(AssistedInjectDagger2Processor())
+        .failsToCompile()
+        .withErrorContaining(
+            "@AssistedModule's @Module must not include AssistedInject_TestModule if @InstallIn is used")
+        .`in`(module).onLine(11)
+  }
+
   @Test fun multipleModulesFails() {
     val moduleOne = JavaFileObjects.forSourceString("test.OneModule", """
       package test;
diff --git a/assisted-inject-processor-dagger2/src/test/java/com/squareup/inject/assisted/dagger2/processor/hiltAnnotations.kt b/assisted-inject-processor-dagger2/src/test/java/com/squareup/inject/assisted/dagger2/processor/hiltAnnotations.kt
new file mode 100644
index 0000000..5424673
--- /dev/null
+++ b/assisted-inject-processor-dagger2/src/test/java/com/squareup/inject/assisted/dagger2/processor/hiltAnnotations.kt
@@ -0,0 +1,19 @@
+package com.squareup.inject.assisted.dagger2.processor
+
+import com.google.testing.compile.JavaFileObjects
+
+val hiltInstallIn = JavaFileObjects.forSourceString("dagger.hilt.InstallIn", """
+  package dagger.hilt;
+
+  import static java.lang.annotation.RetentionPolicy.CLASS;
+  import static java.lang.annotation.ElementType.TYPE;
+
+  import java.lang.annotation.Retention;
+  import java.lang.annotation.Target;
+
+  @Target(TYPE)
+  @Retention(CLASS)
+  public @interface InstallIn {
+    Class<?>[] value();
+  }
+  """.trimIndent())