android gradle plugin task hook

AndroidGradleTaskhook.gradle

import com.android.build.api.transform.Context
import com.android.build.api.transform.DirectoryInput
import com.android.build.api.transform.JarInput
import com.android.build.api.transform.QualifiedContent
import com.android.build.api.transform.SecondaryFile
import com.android.build.api.transform.Transform
import com.android.build.api.transform.TransformException
import com.android.build.api.transform.TransformInput
import com.android.build.api.transform.TransformInvocation
import com.android.build.api.transform.TransformOutputProvider
import com.android.build.gradle.AppPlugin
import com.android.build.gradle.api.ApkVariant
import com.android.build.gradle.api.ApplicationVariant
import com.android.build.gradle.internal.pipeline.TransformTask
import com.android.build.gradle.internal.tasks.PrepareDependenciesTask
import com.android.build.gradle.internal.transforms.DexTransform
import com.android.build.gradle.tasks.MergeSourceSetFolders
import org.gradle.api.Plugin
import org.gradle.api.Project
import org.gradle.api.Task
import org.gradle.api.execution.TaskExecutionListener
import org.gradle.api.plugins.PluginContainer
import org.gradle.api.tasks.TaskState


Map<String, AndroidGradlePluginTaskHooker> taskHookerMap = new HashMap<>()

project.gradle.addListener(new TaskExecutionListener() {
    @Override
    void beforeExecute(Task task) {
        if (task.project == project) {
            if (task in TransformTask) {
                taskHookerMap[task.transform.name]?.beforeExecute(task)
            } else {
                taskHookerMap[task.name]?.beforeExecute(task)
            }
        }

    }

    @Override
    void afterExecute(Task task, TaskState state) {
        if (task.project == project) {
            if (task in TransformTask) {
                taskHookerMap[task.transform.name]?.afterExecute(task)
            } else {
                taskHookerMap[task.name]?.afterExecute(task)
            }
        }
    }
})

project.afterEvaluate {
    project.android.applicationVariants.all { ApplicationVariant appVariant ->
        MergeAssetsHooker mergeAssetsHooker = new MergeAssetsHooker(project, appVariant)
        taskHookerMap.put(mergeAssetsHooker.taskName, mergeAssetsHooker)

        PrepareDependenciesHooker prepareDependenciesHooker = new PrepareDependenciesHooker(project, appVariant)
        taskHookerMap.put(prepareDependenciesHooker.taskName, prepareDependenciesHooker)
        
        DexTransformHooker dexTransformHooker = new DexTransformHooker(project, appVariant)
        taskHookerMap.put(dexTransformHooker.taskName, dexTransformHooker)
    }
}

abstract class AndroidGradlePluginTaskHooker<T extends Task> {

    private Project project

    private ApkVariant apkVariant


    AndroidGradlePluginTaskHooker(Project project, ApkVariant apkVariant) {
        this.project = project
        this.apkVariant = apkVariant
    }

    Project getProject() {
        return this.project
    }


    ApkVariant getApkVariant() {
        return this.apkVariant
    }

    abstract String getTaskName()

    abstract void beforeExecute(T task)

    abstract void afterExecute(T task)
    
    protected void delegate(TransformTask transformTask, TransformWrapper transformWrapper) {
        Field transformFiled = TransformTask.class.getDeclaredField("transform")
        transformFiled.setAccessible(true)
        transformFiled.set(transformTask, transformWrapper)
    }
}

class TransformWrapper extends Transform {

    Transform origin

    TransformWrapper(Transform transform) {
        origin = transform
    }

    @Override
    String getName() {
        return origin.getName()
    }

    @Override
    Set<QualifiedContent.ContentType> getInputTypes() {
        return origin.getInputTypes()
    }

    @Override
    Set<QualifiedContent.ContentType> getOutputTypes() {
        return origin.getOutputTypes()
    }

    @Override
    Set<? super QualifiedContent.Scope> getScopes() {
        return origin.getScopes()
    }

    @Override
    Set<? super QualifiedContent.Scope> getReferencedScopes() {
        return origin.getReferencedScopes()
    }

    @Override
    Collection<File> getSecondaryFileInputs() {
        return origin.getSecondaryFileInputs()
    }

    @Override
    Collection<SecondaryFile> getSecondaryFiles() {
        return origin.getSecondaryFiles()
    }

    @Override
    Collection<File> getSecondaryFileOutputs() {
        return origin.getSecondaryFileOutputs()
    }

    @Override
    Collection<File> getSecondaryDirectoryOutputs() {
        return origin.getSecondaryDirectoryOutputs()
    }

    @Override
    Map<String, Object> getParameterInputs() {
        return origin.getParameterInputs()
    }

    @Override
    boolean isIncremental() {
        return origin.isIncremental()
    }

    @Override
    void transform(Context context, Collection<TransformInput> inputs, Collection<TransformInput> referencedInputs, TransformOutputProvider outputProvider, boolean isIncremental) throws IOException, TransformException, InterruptedException {
        origin.transform(context, inputs, referencedInputs, outputProvider, isIncremental)
    }

    @Override
    void transform(TransformInvocation invocation) throws TransformException, InterruptedException, IOException {
        Collection<TransformInput> inputs = invocation.getInputs()
        for (TransformInput input : inputs) {
            for (DirectoryInput directoryInput : input.getDirectoryInputs()) {
                println "input dir: ${directoryInput.getFile()}"
            }
            for (JarInput jarInput : input.getJarInputs()) {
                println "input jar: ${jarInput.getFile()}"
            }
        }

        Collection<TransformInput> referencedInputs = invocation.getReferencedInputs();
        for (TransformInput transformInput : referencedInputs) {
            for (DirectoryInput directoryInput : transformInput.getDirectoryInputs()) {
                println "referenced input dir: ${directoryInput.getFile()}"
            }
            for (JarInput jarInput : transformInput.getJarInputs()) {
                println "referenced input jar: ${jarInput.getFile()}"
            }
        }

        origin.transform(invocation)
    }

    boolean isCacheable() {
        return origin.isCacheable()
    }

    @Override
    String toString() {
        return origin.toString()
    }

}


class DexTransformHooker extends AndroidGradlePluginTaskHooker<TransformTask> {

    DexTransformHooker(Project project, ApkVariant apkVariant) {
        super(project, apkVariant)
    }

    @Override
    String getTaskName() {
        return "dex"
    }

    @Override
    void beforeExecute(TransformTask task) {
        project.logger.error("hook beforeExecute ${task}")

        delegate(task, new TransformWrapper(task.transform as DexTransform))
    }

    @Override
    void afterExecute(TransformTask task) {
        project.logger.error("hook afterExecute ${task}")
    }
}


class MergeAssetsHooker extends AndroidGradlePluginTaskHooker<MergeSourceSetFolders> {

    MergeAssetsHooker(Project project, ApkVariant apkVariant) {
        super(project, apkVariant)
    }

    @Override
    String getTaskName() {
        return "merge${apkVariant.name.capitalize()}Assets"
    }

    @Override
    void beforeExecute(MergeSourceSetFolders task) {
        project.logger.error("hook beforeTaskExecute ${task}")

    }

    @Override
    void afterExecute(MergeSourceSetFolders task) {
        project.logger.error("hook afterTaskExecute ${task}")
    }
}

class PrepareDependenciesHooker extends AndroidGradlePluginTaskHooker<PrepareDependenciesTask> {

    PrepareDependenciesHooker(Project project, ApkVariant apkVariant) {
        super(project, apkVariant)
    }

    @Override
    String getTaskName() {
        return "prepare${apkVariant.name.capitalize()}Dependencies"
    }

    @Override
    void beforeExecute(PrepareDependenciesTask task) {
        project.logger.error("hook beforeTaskExecute ${task}")
    }

    @Override
    void afterExecute(PrepareDependenciesTask task) {
        project.logger.error("hook afterTaskExecute ${task}")
    }

}