build.gradle

apply plugin: 'com.android.library'
apply plugin: 'maven'

android {
    compileSdkVersion rootProject.compileSdkVersion
    buildToolsVersion rootProject.buildToolsVersion

    defaultConfig {
        minSdkVersion rootProject.minSdkVersion
        targetSdkVersion rootProject.targetSdkVersion
        versionCode 0
        versionName "0.1"

        testInstrumentationRunner "androidx.test.runner.AndroidJUnitRunner"
        ndk {
            abiFilters ABI_FILTERS.split(",")
        }
        externalNativeBuild {
            cmake {
              if(System.env.BUILD_LITE_INTERPRETER == '0') {
                arguments "-DANDROID_STL=c++_shared", "-DBUILD_LITE_INTERPRETER=OFF", "-DUSE_LITE_INTERPRETER_PROFILER=OFF"
              } else {
                arguments "-DANDROID_STL=c++_shared", "-DUSE_LITE_INTERPRETER_PROFILER=OFF"
              }
            }
        }
    }
    buildTypes {
        debug {
            minifyEnabled false
            debuggable true
        }
        release {
            minifyEnabled false
        }
    }
    sourceSets {
        main {
            java {
              if(System.env.BUILD_LITE_INTERPRETER == '0') {
                println 'Build pytorch_jni'
                exclude 'org/pytorch/LiteModuleLoader.java'
                exclude 'org/pytorch/LiteNativePeer.java'
              } else {
                println 'Build pytorch_jni_lite'
              }
            }
            jniLibs.srcDirs = ['src/main/jniLibs']
            manifest.srcFile 'src/main/AndroidManifest.xml'
        }
        androidTest {
            java {
                if(System.env.BUILD_LITE_INTERPRETER == '0') {
                    println 'Build test for full jit (pytorch_jni)'
                    exclude 'org/pytorch/PytorchHostTests.java'
                    exclude 'org/pytorch/PytorchLiteInstrumentedTests.java'
                    exclude 'org/pytorch/suite/PytorchLiteInstrumentedTestSuite.java'
                } else {
                    println 'Build test for lite interpreter (pytorch_jni_lite)'
                    exclude 'org/pytorch/PytorchHostTests.java'
                    exclude 'org/pytorch/PytorchInstrumentedTests.java'
                    exclude 'org/pytorch/suite/PytorchInstrumentedTestSuite.java'
                }
            }
        }
    }
    externalNativeBuild {
        cmake {
            path "CMakeLists.txt"
        }
    }

    packagingOptions {
        if (nativeLibsDoNotStrip.toBoolean()) {
            doNotStrip "**/*.so"
            logger.warn('WARNING: nativeLibsDoNotStrip==true; debug symbols included')
        }
    }

    useLibrary 'android.test.runner'
    useLibrary 'android.test.base'
    useLibrary 'android.test.mock'
}

dependencies {
    implementation 'com.facebook.fbjni:fbjni-java-only:' + rootProject.fbjniJavaOnlyVersion
    implementation 'com.facebook.soloader:nativeloader:' + rootProject.soLoaderNativeLoaderVersion

    testImplementation 'junit:junit:' + rootProject.junitVersion
    testImplementation 'androidx.test:core:' + rootProject.coreVersion

    androidTestImplementation 'junit:junit:' + rootProject.junitVersion
    androidTestImplementation 'androidx.test:core:' + rootProject.coreVersion
    androidTestImplementation 'androidx.test.ext:junit:' + rootProject.extJUnitVersion
    androidTestImplementation 'androidx.test:rules:' + rootProject.rulesVersion
    androidTestImplementation 'androidx.test:runner:' + rootProject.runnerVersion
}

apply from: rootProject.file('gradle/release.gradle')

task sourcesJar(type: Jar) {
    from android.sourceSets.main.java.srcDirs
    classifier = 'sources'
}

def getLibtorchHeadersDir() {
  def abi = ABI_FILTERS.split(",")[0]
  return "$rootDir/pytorch_android/src/main/cpp/libtorch_include/$abi"
}

afterEvaluate {
  if (POM_PACKAGING == 'aar') {
    android.libraryVariants.all { variant ->
      variant.outputs.each { output ->
        File f = output.outputFile
        if (f.name.endsWith(".aar")) {
          output.assemble.finalizedBy addFolderToAarTask(
              "addHeadersToAar" + variant.name,
              f.path,
              getLibtorchHeadersDir(),
              "headers")
        }
      }
    }
  }
}

tasks.whenTaskAdded { task ->
  if (task.name.startsWith("bundle") && task.name.endsWith("Aar")) {
    doLast {
      addFolderToAar("addHeadersTo" + task.name, task.archivePath, getLibtorchHeadersDir(), 'headers')
    }
  }
}

def addFolderToAarTask(taskName, aarPath, folderPath, folderPathInAar) {
  return tasks.register(taskName) {
    doLast {
      addFolderToAar(taskName, aarPath, folderPath, folderPathInAar)
    }
  }
}

def addFolderToAar(taskName, aarPath, folderPath, folderPathInAar) {
    def tmpDir = file("${buildDir}/${taskName}")
    tmpDir.mkdir()
    def tmpDirFolder = file("${tmpDir.path}/${folderPathInAar}")
    tmpDirFolder.mkdir()
    copy {
      from zipTree(aarPath)
      into tmpDir
    }
    copy {
      from fileTree(folderPath)
      into tmpDirFolder
    }
    ant.zip(destfile: aarPath) {
      fileset(dir: tmpDir.path)
    }
    delete tmpDir
}

artifacts.add('archives', sourcesJar)