Anvil Code Generator

ContributesApiCodeGenerator.kt

package com.tonal.trainer.anvilcompilers

import com.google.auto.service.AutoService
import com.squareup.anvil.annotations.ContributesTo
import com.squareup.anvil.compiler.api.AnvilContext
import com.squareup.anvil.compiler.api.CodeGenerator
import com.squareup.anvil.compiler.api.GeneratedFile
import com.squareup.anvil.compiler.api.createGeneratedFile
import com.squareup.anvil.compiler.internal.asClassName
import com.squareup.anvil.compiler.internal.buildFile
import com.squareup.anvil.compiler.internal.classesAndInnerClass
import com.squareup.anvil.compiler.internal.fqName
import com.squareup.anvil.compiler.internal.hasAnnotation
import com.squareup.kotlinpoet.AnnotationSpec
import com.squareup.kotlinpoet.ClassName
import com.squareup.kotlinpoet.FileSpec
import com.squareup.kotlinpoet.FunSpec
import com.squareup.kotlinpoet.ParameterSpec
import com.squareup.kotlinpoet.TypeSpec
import com.squareup.kotlinpoet.asClassName
import com.tonal.trainer.anvilannotations.ContributesNonUserApi
import com.tonal.trainer.anvilannotations.ContributesUserApi
import com.tonal.trainer.lib.daggerscopes.AppComponent
import com.tonal.trainer.lib.daggerscopes.UserComponent
import dagger.Module
import dagger.Provides
import dagger.Reusable
import org.jetbrains.kotlin.descriptors.ModuleDescriptor
import org.jetbrains.kotlin.name.FqName
import org.jetbrains.kotlin.psi.KtClassOrObject
import org.jetbrains.kotlin.psi.KtFile
import java.io.File

/**
 * This Anvil code generator allows you to annotate retrofit interfaces with @ContributesUserApi
 * or @ContributesNonUserApi. Doing so will automatically generate a Dagger module that provides
 * the retrofit interface.
 * 
 * @ContributesUserApi interfaces will generate a module that injects the @UserApi tagged Retrofit
 * instance which has the user auth header interceptor and the userId path interceptor which replaces
 * {userId} in url paths with the actual userId automatically.
 */
@AutoService(CodeGenerator::class)
class ContributesApiCodeGenerator : CodeGenerator {

    override fun isApplicable(context: AnvilContext): Boolean = true

    override fun generateCode(codeGenDir: File, module: ModuleDescriptor, projectFiles: Collection<KtFile>): Collection<GeneratedFile> {
        return projectFiles.classesAndInnerClass(module)
            .mapNotNull { clazz ->
                when {
                    clazz.hasAnnotation(ContributesUserApi::class.fqName, module) -> generateModule(clazz, isUserApi = true, codeGenDir, module)
                    clazz.hasAnnotation(ContributesNonUserApi::class.fqName, module) -> generateModule(clazz, isUserApi = false, codeGenDir, module)
                    else -> null
                }
            }
            .toList()
    }

    private fun generateModule(apiClass: KtClassOrObject, isUserApi: Boolean, codeGenDir: File, module: ModuleDescriptor): GeneratedFile {
        val generatedPackage = apiClass.containingKtFile.packageFqName.toString()
        val moduleClassName = "${apiClass.name}_Module"
        val component = if (isUserApi) UserComponent::class.asClassName() else AppComponent::class.asClassName()
        val content = FileSpec.buildFile(generatedPackage, moduleClassName) {
            addType(
                TypeSpec.classBuilder(moduleClassName)
                    .addAnnotation(Module::class)
                    .addAnnotation(AnnotationSpec.builder(ContributesTo::class).addMember("%T::class", component).build())
                    .addFunction(
                        FunSpec.builder("provide${apiClass.name}")
                            .addParameter(
                                ParameterSpec.builder("retrofit", ClassName("retrofit2", "Retrofit"))
                                    .apply {
                                        if (isUserApi) {
                                            // If this is a user api, inject `@UserApi Retrofit` instead of `Retrofit`
                                            addAnnotation(userApiFqName.asClassName(module))
                                        }
                                    }
                                    .build(),
                            )
                            .returns(apiClass.asClassName())
                            .addAnnotation(Provides::class)
                            .addAnnotation(Reusable::class)
                            .addCode("return retrofit.create(%T::class.java)", apiClass.asClassName())
                            .build(),
                    )
                    .build(),
            )
        }
        return createGeneratedFile(codeGenDir, generatedPackage, moduleClassName, content)
    }

    companion object {
        private val userApiFqName = FqName("com.tonal.trainer.lib.data.db.UserApi")
    }
}

ContributesViewModelCodeGenerator.kt

package com.tonal.trainer.anvilcompilers

import com.google.auto.service.AutoService
import com.squareup.anvil.annotations.ContributesTo
import com.squareup.anvil.compiler.api.AnvilCompilationException
import com.squareup.anvil.compiler.api.AnvilContext
import com.squareup.anvil.compiler.api.CodeGenerator
import com.squareup.anvil.compiler.api.GeneratedFile
import com.squareup.anvil.compiler.api.createGeneratedFile
import com.squareup.anvil.compiler.internal.asClassName
import com.squareup.anvil.compiler.internal.buildFile
import com.squareup.anvil.compiler.internal.classesAndInnerClass
import com.squareup.anvil.compiler.internal.fqName
import com.squareup.anvil.compiler.internal.hasAnnotation
import com.squareup.anvil.compiler.internal.requireFqName
import com.squareup.anvil.compiler.internal.requireTypeReference
import com.squareup.anvil.compiler.internal.scope
import com.squareup.kotlinpoet.AnnotationSpec
import com.squareup.kotlinpoet.ClassName
import com.squareup.kotlinpoet.FileSpec
import com.squareup.kotlinpoet.FunSpec
import com.squareup.kotlinpoet.KModifier
import com.squareup.kotlinpoet.ParameterizedTypeName.Companion.parameterizedBy
import com.squareup.kotlinpoet.STAR
import com.squareup.kotlinpoet.TypeSpec
import com.tonal.trainer.anvilannotations.ContributesViewModel
import dagger.Binds
import dagger.Module
import dagger.assisted.Assisted
import dagger.assisted.AssistedFactory
import dagger.assisted.AssistedInject
import dagger.multibindings.IntoMap
import org.jetbrains.kotlin.descriptors.ModuleDescriptor
import org.jetbrains.kotlin.name.FqName
import org.jetbrains.kotlin.psi.KtClassOrObject
import org.jetbrains.kotlin.psi.KtFile
import org.jetbrains.kotlin.psi.allConstructors
import java.io.File

/**
 * This Anvil code generator allows you to @AssistedInject a ViewModel without registering it in a Dagger
 * Module by hand.
 */
@AutoService(CodeGenerator::class)
class ContributesViewModelCodeGenerator : CodeGenerator {

    override fun isApplicable(context: AnvilContext): Boolean = true

    override fun generateCode(codeGenDir: File, module: ModuleDescriptor, projectFiles: Collection<KtFile>): Collection<GeneratedFile> {
        return projectFiles.classesAndInnerClass(module)
            .filter { it.hasAnnotation(ContributesViewModel::class.fqName, module) }
            .flatMap { listOf(generateModule(it, codeGenDir, module), generateAssistedFactory(it, codeGenDir, module)) }
            .toList()
    }

    private fun generateModule(vmClass: KtClassOrObject, codeGenDir: File, module: ModuleDescriptor): GeneratedFile {
        val generatedPackage = vmClass.containingKtFile.packageFqName.toString()
        val moduleClassName = "${vmClass.name}_Module"
        val scope = vmClass.scope(ContributesViewModel::class.fqName, module)
        val content = FileSpec.buildFile(generatedPackage, moduleClassName) {
            addType(
                TypeSpec.classBuilder(moduleClassName)
                    .addModifiers(KModifier.ABSTRACT)
                    .addAnnotation(Module::class)
                    .addAnnotation(AnnotationSpec.builder(ContributesTo::class).addMember("%T::class", scope.asClassName(module)).build())
                    .addFunction(
                        FunSpec.builder("bind${vmClass.name}Factory")
                            .addModifiers(KModifier.ABSTRACT)
                            .addParameter("factory", ClassName(generatedPackage, "${vmClass.name}_AssistedFactory"))
                            .returns(tonalViewModelFactoryFqName.asClassName(module).parameterizedBy(STAR, STAR))
                            .addAnnotation(Binds::class)
                            .addAnnotation(IntoMap::class)
                            .addAnnotation(AnnotationSpec.builder(viewModelKeyFqName.asClassName(module)).addMember("%T::class", vmClass.asClassName()).build())
                            .build(),
                    )
                    .build(),
            )
        }
        return createGeneratedFile(codeGenDir, generatedPackage, moduleClassName, content)
    }

    private fun generateAssistedFactory(vmClass: KtClassOrObject, codeGenDir: File, module: ModuleDescriptor): GeneratedFile {
        val generatedPackage = vmClass.containingKtFile.packageFqName.toString()
        val assistedFactoryClassName = "${vmClass.name}_AssistedFactory"
        val constructor = vmClass.allConstructors.singleOrNull { it.hasAnnotation(AssistedInject::class.fqName, module) }
        val assistedParameter = constructor?.valueParameters?.singleOrNull { it.hasAnnotation(Assisted::class.fqName, module) }
        if (constructor == null || assistedParameter == null) {
            throw AnvilCompilationException(
                "${vmClass.requireFqName()} must have an @AssistedInject constructor with @Assisted initialState: S parameter",
                element = vmClass.identifyingElement,
            )
        }
        if (assistedParameter.name != "initialState") {
            throw AnvilCompilationException(
                "${vmClass.requireFqName()} @Assisted parameter must be named initialState",
                element = assistedParameter.identifyingElement,
            )
        }
        val vmClassName = vmClass.asClassName()
        val stateClassName = assistedParameter.requireTypeReference(module).requireFqName(module).asClassName(module)
        val content = FileSpec.buildFile(generatedPackage, assistedFactoryClassName) {
            addType(
                TypeSpec.interfaceBuilder(assistedFactoryClassName)
                    .addSuperinterface(tonalViewModelFactoryFqName.asClassName(module).parameterizedBy(vmClassName, stateClassName))
                    .addAnnotation(AssistedFactory::class)
                    .addFunction(
                        FunSpec.builder("create")
                            .addModifiers(KModifier.OVERRIDE, KModifier.ABSTRACT)
                            .addParameter("initialState", stateClassName)
                            .returns(vmClassName)
                            .build(),
                    )
                    .build(),
            )
        }
        return createGeneratedFile(codeGenDir, generatedPackage, assistedFactoryClassName, content)
    }

    companion object {
        private val tonalViewModelFactoryFqName = FqName("com.tonal.trainer.lib.ui.TonalViewModelFactory")
        private val viewModelKeyFqName = FqName("com.tonal.trainer.lib.ui.ViewModelKey")
    }
}

MyViewModel.kt

/**
 * Example ViewModel that uses the @ContributesViewModel code generator.
 */
@ContributesViewModel(UserComponent::class)
class MyViewModel @AssistedInject constructor(
    @Assisted initialState: MyState,
    myRepository: MyRepository,
    userProvider: UserProvider,
): MavericksViewModel<MyState>(initialState) {
    companion object: MavericksViewModelFactory<MyViewModel, MyState> by daggerMavericksViewModelFactory()
}

TonalViewModelFactory.kt

package com.tonal.trainer.lib.ui

import com.airbnb.mvrx.MavericksState

/**
 * Helper interface used to make using AssistedInject easier.
 */
interface TonalViewModelFactory<VM : TonalViewModel<S>, S : MavericksState> {
    fun create(initialState: S): VM
}

ViewModelKey.kt

package com.tonal.trainer.lib.ui

import dagger.MapKey
import kotlin.reflect.KClass

@Target(AnnotationTarget.FUNCTION, AnnotationTarget.PROPERTY_GETTER, AnnotationTarget.PROPERTY_SETTER)
@Retention(AnnotationRetention.RUNTIME)
@MapKey
annotation class ViewModelKey(val value: KClass<out TonalViewModel<*>>)