SBT 0.10 thrift generation plugin

CompileThrift.scala

package com.twitter.sbt

import sbt._
import sbt.Keys._
import sbt.Project.Initialize
import sbt.Process._
import scala.collection.JavaConversions._
import java.io.File

object CompileThrift extends Plugin {
  val env = System.getenv().toMap
  val thriftBin = SettingKey[File]("thrift-bin", "Thrift generator")
  val thriftSources = SettingKey[File]("thrift-sources", "A directory containing thrift files")
  val thriftAutoCompile = SettingKey[Boolean]("thrift-auto-compile", "Auto compile thrift or no")

  private def compileChanged (sources: File, target: File) = {
    val minSourceMod = sources.listFiles.foldLeft(java.lang.Long.MAX_VALUE) { (last, f) =>
      f.lastModified min last
    }
    val maxTargetMod = target.listFiles.foldLeft(0L) { (last, f) =>
      f.lastModified max last
    }
    minSourceMod > maxTargetMod
  }

  def compileDir(thriftBin: File, thriftSources: File, targetDir: File, lang: String, out: Logger): Seq[File] = {
    val outputDir = targetDir / ("gen-" + lang)
    Process("mkdir -p " + outputDir).run
    if (compileChanged(thriftSources, outputDir)) {
      out.info("generating thrift files")
      thriftSources.listFiles.map {file =>
        compile(thriftBin, file, targetDir, lang, out)
      }
      val generatedFiles = outputDir ** "*.java"
      val rv = generatedFiles.get
      out.debug("generated the following files:\n\t" + rv.mkString("\n\t"))
      rv
    } else {
      out.info("nothing to do for thrift files")
      Nil
    }
  }

  def compile(bin: File, thriftSource: File, targetDir: File, lang: String, out: Logger) {
    val cmd = "%s -gen %s -o %s %s".format(bin, lang, targetDir.absolutePath, thriftSource)
    out.debug("generating thrift with " + cmd)
    Process(cmd).run.exitValue
  }

  def thriftCompileTask(lang: String, check: TaskKey[Boolean]) = {
    (streams, check, thriftBin, thriftSources, (sourceManaged in Compile)) map {
      (out, check, bin, source, target) =>
        if (check) {
          compileDir(bin, source, target, lang, out.log)
        } else {
          Nil
        }
    }
  }

  type ThriftKeySet = (SettingKey[Boolean], TaskKey[Boolean], TaskKey[Seq[File]])

  def langToKeyName(lang: String) = "thrift-" + lang

  def makeLangKeysFor(langs: String*) = {
    Map(langs.map { lang =>
      lang -> makeLangKeys(lang)
    } : _*)
  }

  def makeLangKeys(lang: String): ThriftKeySet = {
    val thriftLang = langToKeyName(lang)
    val enableFlag = SettingKey[Boolean](thriftLang + "-enable", "Flag to enable generating " + lang)
    val enableTask = TaskKey[Boolean](thriftLang + "-check", "Should I compile " + lang)
    val compileTask = TaskKey[Seq[File]](thriftLang + "-compile", "compile thrift " + lang)
    (enableFlag, enableTask, compileTask)
  }

  val keysMap = makeLangKeysFor("java", "python", "rb")
  val (thriftJavaEnable, thriftJavaCheck, thriftJavaCompile) = keysMap("java")
  val (thriftPythonEnable, thriftPythonCheck, thriftPythonCompile) = keysMap("python")
  val (thriftRubyEnable, thriftRubyCheck, thriftRubyCompile) = keysMap("rb")

  def makeLangTasksFor(langs: String*): Seq[Setting[_]] = {
    langs.map {lang =>
      makeLangTasks(lang)
    }.flatten
  }

  def makeLangTasks(lang: String): Seq[Setting[_]] = {
    val keySet = keysMap(lang)
    Seq (keySet._1 := true,
         keySet._2 <<= keySet._1 map identity,
         keySet._3 <<= thriftCompileTask(lang, keySet._2),
         sourceGenerators in Compile <+= thriftCompileTask(lang, keySet._2))
  }

  def thriftSettings(langs: String*) = {
    Seq(
      thriftBin := new File(env.get("THRIFT_BIN").getOrElse("thrift")),
      thriftSources <<= (sourceDirectory in Compile) { _ / "thrift" },
      thriftAutoCompile := true
    ) ++ makeLangTasksFor(langs: _*)
  }
}