1 files changed, 83 insertions, 1 deletions

diff --git a/project/SparkBuild.scala b/project/SparkBuild.scala
index 9e6d341ca2..ffbca25e46 100644
--- a/project/SparkBuild.scala
+++ b/project/SparkBuild.scala
@@ -18,6 +18,7 @@
 import java.io._
 import java.nio.file.Files
 
+import scala.io.Source
 import scala.util.Properties
 import scala.collection.JavaConverters._
 import scala.collection.mutable.Stack
@@ -29,6 +30,8 @@ import sbtunidoc.Plugin.UnidocKeys.unidocGenjavadocVersion
 import com.simplytyped.Antlr4Plugin._
 import com.typesafe.sbt.pom.{PomBuild, SbtPomKeys}
 import com.typesafe.tools.mima.plugin.MimaKeys
+import org.scalastyle.sbt.ScalastylePlugin._
+import org.scalastyle.sbt.Tasks
 
 import spray.revolver.RevolverPlugin._
 
@@ -147,7 +150,86 @@ object SparkBuild extends PomBuild {
       "org.spark-project" %% "genjavadoc-plugin" % unidocGenjavadocVersion.value cross CrossVersion.full),
     scalacOptions <+= target.map(t => "-P:genjavadoc:out=" + (t / "java")))
 
-  lazy val sharedSettings = sparkGenjavadocSettings ++ Seq (
+  lazy val scalaStyleRules = Project("scalaStyleRules", file("scalastyle"))
+    .settings(
+      libraryDependencies += "org.scalastyle" %% "scalastyle" % "0.8.0"
+    )
+
+  lazy val scalaStyleOnCompile = taskKey[Unit]("scalaStyleOnCompile")
+
+  lazy val scalaStyleOnTest = taskKey[Unit]("scalaStyleOnTest")
+
+  // We special case the 'println' lint rule to only be a warning on compile, because adding
+  // printlns for debugging is a common use case and is easy to remember to remove.
+  val scalaStyleOnCompileConfig: String = {
+    val in = "scalastyle-config.xml"
+    val out = "scalastyle-on-compile.generated.xml"
+    val replacements = Map(
+      """customId="println" level="error"""" -> """customId="println" level="warn""""
+    )
+    var contents = Source.fromFile(in).getLines.mkString("\n")
+    for ((k, v) <- replacements) {
+      require(contents.contains(k), s"Could not rewrite '$k' in original scalastyle config.")
+      contents = contents.replace(k, v)
+    }
+    new PrintWriter(out) {
+      write(contents)
+      close()
+    }
+    out
+  }
+
+  // Return a cached scalastyle task for a given configuration (usually Compile or Test)
+  private def cachedScalaStyle(config: Configuration) = Def.task {
+    val logger = streams.value.log
+    // We need a different cache dir per Configuration, otherwise they collide
+    val cacheDir = target.value / s"scalastyle-cache-${config.name}"
+    val cachedFun = FileFunction.cached(cacheDir, FilesInfo.lastModified, FilesInfo.exists) {
+      (inFiles: Set[File]) => {
+        val args: Seq[String] = Seq.empty
+        val scalaSourceV = Seq(file(scalaSource.in(config).value.getAbsolutePath))
+        val configV = (baseDirectory in ThisBuild).value / scalaStyleOnCompileConfig
+        val configUrlV = scalastyleConfigUrl.in(config).value
+        val streamsV = streams.in(config).value
+        val failOnErrorV = true
+        val scalastyleTargetV = scalastyleTarget.in(config).value
+        val configRefreshHoursV = scalastyleConfigRefreshHours.in(config).value
+        val targetV = target.in(config).value
+        val configCacheFileV = scalastyleConfigUrlCacheFile.in(config).value
+
+        logger.info(s"Running scalastyle on ${name.value} in ${config.name}")
+        Tasks.doScalastyle(args, configV, configUrlV, failOnErrorV, scalaSourceV, scalastyleTargetV,
+          streamsV, configRefreshHoursV, targetV, configCacheFileV)
+
+        Set.empty
+      }
+    }
+
+    cachedFun(findFiles(scalaSource.in(config).value))
+  }
+
+  private def findFiles(file: File): Set[File] = file.isDirectory match {
+    case true => file.listFiles().toSet.flatMap(findFiles) + file
+    case false => Set(file)
+  }
+
+  def enableScalaStyle: Seq[sbt.Def.Setting[_]] = Seq(
+    scalaStyleOnCompile := cachedScalaStyle(Compile).value,
+    scalaStyleOnTest := cachedScalaStyle(Test).value,
+    logLevel in scalaStyleOnCompile := Level.Warn,
+    logLevel in scalaStyleOnTest := Level.Warn,
+    (compile in Compile) := {
+      scalaStyleOnCompile.value
+      (compile in Compile).value
+    },
+    (compile in Test) := {
+      scalaStyleOnTest.value
+      (compile in Test).value
+    }
+  )
+
+  lazy val sharedSettings = sparkGenjavadocSettings ++
+      (if (sys.env.contains("NOLINT_ON_COMPILE")) Nil else enableScalaStyle) ++ Seq(
     exportJars in Compile := true,
     exportJars in Test := false,
     javaHome := sys.env.get("JAVA_HOME")