fix #233 add append and exclude rules to assembly

build.sc

@@ -1,6 +1,5 @@
 import $file.shared
 import $file.upload
-import java.io.File
 import java.nio.file.attribute.PosixFilePermission
 
 import ammonite.ops._
@@ -9,7 +8,6 @@ import mill._
 import mill.scalalib._
 import publish._
 import mill.modules.Jvm.createAssembly
-
 import upickle.Js
 trait MillPublishModule extends PublishModule{
 

main/src/mill/modules/Assembly.scala

@@ -0,0 +1,127 @@
+package mill.modules
+
+import java.io.InputStream
+import java.util.jar.JarFile
+import java.util.regex.Pattern
+
+import ammonite.ops._
+import geny.Generator
+import mill.Agg
+
+import scala.collection.JavaConverters._
+
+object Assembly {
+
+  val defaultRules: Seq[Rule] = Seq(
+    Rule.Append("reference.conf"),
+    Rule.Exclude(JarFile.MANIFEST_NAME),
+    Rule.ExcludePattern(".*\\.[sS][fF]"),
+    Rule.ExcludePattern(".*\\.[dD][sS][aA]"),
+    Rule.ExcludePattern(".*\\.[rR][sS][aA]")
+  )
+
+  sealed trait Rule extends Product with Serializable
+  object Rule {
+    case class Append(path: String) extends Rule
+
+    object AppendPattern {
+      def apply(pattern: String): AppendPattern = AppendPattern(Pattern.compile(pattern))
+    }
+    case class AppendPattern(pattern: Pattern) extends Rule
+
+    case class Exclude(path: String) extends Rule
+
+    object ExcludePattern {
+      def apply(pattern: String): ExcludePattern = ExcludePattern(Pattern.compile(pattern))
+    }
+    case class ExcludePattern(pattern: Pattern) extends Rule
+  }
+
+  def groupAssemblyEntries(inputPaths: Agg[Path], assemblyRules: Seq[Assembly.Rule]): Map[String, GroupedEntry] = {
+    val rulesMap = assemblyRules.collect {
+      case r@Rule.Append(path) => path -> r
+      case r@Rule.Exclude(path) => path -> r
+    }.toMap
+
+    val appendPatterns = assemblyRules.collect {
+      case Rule.AppendPattern(pattern) => pattern.asPredicate().test(_)
+    }
+
+    val excludePatterns = assemblyRules.collect {
+      case Rule.ExcludePattern(pattern) => pattern.asPredicate().test(_)
+    }
+
+    classpathIterator(inputPaths).foldLeft(Map.empty[String, GroupedEntry]) {
+      case (entries, entry) =>
+        val mapping = entry.mapping
+
+        rulesMap.get(mapping) match {
+          case Some(_: Assembly.Rule.Exclude) =>
+            entries
+          case Some(_: Assembly.Rule.Append) =>
+            val newEntry = entries.getOrElse(mapping, AppendEntry.empty).append(entry)
+            entries + (mapping -> newEntry)
+
+          case _ if excludePatterns.exists(_(mapping)) =>
+            entries
+          case _ if appendPatterns.exists(_(mapping)) =>
+            val newEntry = entries.getOrElse(mapping, AppendEntry.empty).append(entry)
+            entries + (mapping -> newEntry)
+
+          case _ if !entries.contains(mapping) =>
+            entries + (mapping -> WriteOnceEntry(entry))
+          case _ =>
+            entries
+        }
+    }
+  }
+
+  private def classpathIterator(inputPaths: Agg[Path]): Generator[AssemblyEntry] = {
+    Generator.from(inputPaths)
+      .filter(exists)
+      .flatMap {
+        p =>
+          if (p.isFile) {
+            val jf = new JarFile(p.toIO)
+            Generator.from(
+              for(entry <- jf.entries().asScala if !entry.isDirectory)
+                yield JarFileEntry(entry.getName, () => jf.getInputStream(entry))
+            )
+          }
+          else {
+            ls.rec.iter(p)
+              .filter(_.isFile)
+              .map(sub => PathEntry(sub.relativeTo(p).toString, sub))
+          }
+      }
+  }
+}
+
+private[modules] sealed trait GroupedEntry {
+  def append(entry: AssemblyEntry): GroupedEntry
+}
+
+private[modules] object AppendEntry {
+  val empty: AppendEntry = AppendEntry(Nil)
+}
+
+private[modules] case class AppendEntry(entries: List[AssemblyEntry]) extends GroupedEntry {
+  def append(entry: AssemblyEntry): GroupedEntry = copy(entries = entry :: this.entries)
+}
+
+private[modules] case class WriteOnceEntry(entry: AssemblyEntry) extends GroupedEntry {
+  def append(entry: AssemblyEntry): GroupedEntry = this
+}
+
+private[this] sealed trait AssemblyEntry {
+  def mapping: String
+  def inputStream: InputStream
+}
+
+private[this] case class PathEntry(mapping: String, path: Path) extends AssemblyEntry {
+  def inputStream: InputStream = read.getInputStream(path)
+}
+
+private[this] case class JarFileEntry(mapping: String, getIs: () => InputStream) extends AssemblyEntry {
+  def inputStream: InputStream = getIs()
+}

main/src/mill/modules/Jvm.scala

@@ -1,14 +1,14 @@
 package mill.modules
 
-import java.io.{ByteArrayInputStream, File, FileOutputStream}
+import java.io._
 import java.lang.reflect.Modifier
-import java.net.{URI, URLClassLoader}
-import java.nio.file.{FileSystems, Files, OpenOption, StandardOpenOption}
+import java.net.URI
+import java.nio.file.{FileSystems, Files, StandardOpenOption}
 import java.nio.file.attribute.PosixFilePermission
+import java.util.Collections
 import java.util.jar.{JarEntry, JarFile, JarOutputStream}
 
 import ammonite.ops._
-import ammonite.util.Util
 import coursier.{Cache, Dependency, Fetch, Repository, Resolution}
 import geny.Generator
 import mill.main.client.InputPumper
@@ -17,7 +17,7 @@ import mill.util.{Ctx, IO}
 import mill.util.Loose.Agg
 
 import scala.collection.mutable
-
+import scala.collection.JavaConverters._
 
 object Jvm {
 
@@ -232,18 +232,20 @@ object Jvm {
     PathRef(outputPath)
   }
 
-  def newOutputStream(p: java.nio.file.Path) = Files.newOutputStream(
-    p,
-    StandardOpenOption.TRUNCATE_EXISTING,
-    StandardOpenOption.CREATE
-  )
+  def newOutputStream(p: java.nio.file.Path, append: Boolean = false) = {
+    val options =
+      if(append) Seq(StandardOpenOption.APPEND)
+      else Seq(StandardOpenOption.TRUNCATE_EXISTING, StandardOpenOption.CREATE)
+    Files.newOutputStream(p, options:_*)
+  }
 
   def createAssembly(inputPaths: Agg[Path],
                      mainClass: Option[String] = None,
                      prependShellScript: String = "",
                      base: Option[Path] = None,
-                     isWin: Boolean = scala.util.Properties.isWin)
-                    (implicit ctx: Ctx.Dest) = {
+                     assemblyRules: Seq[Assembly.Rule] = Assembly.defaultRules)
+                    (implicit ctx: Ctx.Dest with Ctx.Log): PathRef = {
+
     val tmp = ctx.dest / "out-tmp.jar"
 
     val baseUri = "jar:" + tmp.toIO.getCanonicalFile.toURI.toASCIIString
@@ -263,20 +265,22 @@ object Jvm {
     manifest.write(manifestOut)
     manifestOut.close()
 
-    def isSignatureFile(mapping: String): Boolean =
-      Set("sf", "rsa", "dsa").exists(ext => mapping.toLowerCase.endsWith(s".$ext"))
-
-    for(v <- classpathIterator(inputPaths)){
-      val (file, mapping) = v
-      val p = zipFs.getPath(mapping)
-      if (p.getParent != null) Files.createDirectories(p.getParent)
-      if (!isSignatureFile(mapping)) {
-        val outputStream = newOutputStream(p)
-        IO.stream(file, outputStream)
-        outputStream.close()
+    Assembly.groupAssemblyEntries(inputPaths, assemblyRules).view
+      .map {
+        case (mapping, aggregate) =>
+          zipFs.getPath(mapping) -> aggregate
       }
-      file.close()
-    }
+      .foreach {
+        case (path, AppendEntry(entries)) =>
+          val concatenated = new SequenceInputStream(
+            Collections.enumeration(entries.map(_.inputStream).asJava))
+          writeEntry(path, concatenated, append = Files.exists(path))
+        case (path, WriteOnceEntry(entry)) =>
+          if (Files.notExists(path)) {
+            writeEntry(path, entry.inputStream, append = false)
+          }
+      }
+
     zipFs.close()
     val output = ctx.dest / "out.jar"
 
@@ -301,28 +305,14 @@ object Jvm {
     PathRef(output)
   }
 
+  private def writeEntry(p: java.nio.file.Path, is: InputStream, append: Boolean): Unit = {
+    if (p.getParent != null) Files.createDirectories(p.getParent)
+    val outputStream = newOutputStream(p, append)
 
-  def classpathIterator(inputPaths: Agg[Path]) = {
-    Generator.from(inputPaths)
-      .filter(exists)
-      .flatMap{
-        p =>
-          if (p.isFile) {
-            val jf = new JarFile(p.toIO)
-            import collection.JavaConverters._
-            Generator.selfClosing((
-              for(entry <- jf.entries().asScala if !entry.isDirectory)
-                yield (jf.getInputStream(entry), entry.getName),
-              () => jf.close()
-            ))
-          }
-          else {
-            ls.rec.iter(p)
-              .filter(_.isFile)
-              .map(sub => read.getInputStream(sub) -> sub.relativeTo(p).toString)
-          }
-      }
+    IO.stream(is, outputStream)
 
+    outputStream.close()
+    is.close()
   }
 
   def universalScript(shellCommands: String,

scalalib/src/mill/scalalib/JavaModule.scala

@@ -7,7 +7,7 @@ import coursier.Repository
 import mill.define.Task
 import mill.define.TaskModule
 import mill.eval.{PathRef, Result}
-import mill.modules.Jvm
+import mill.modules.{Assembly, Jvm}
 import mill.modules.Jvm.{createAssembly, createJar}
 import Lib._
 import mill.scalalib.publish.{Artifact, Scope}
@@ -108,6 +108,8 @@ trait JavaModule extends mill.Module with TaskModule { outer =>
     }
   }
 
+  def assemblyRules: Seq[Assembly.Rule] = Assembly.defaultRules
+
   def sources = T.sources{ millSourcePath / 'src }
   def resources = T.sources{ millSourcePath / 'resources }
   def generatedSources = T{ Seq.empty[PathRef] }
@@ -130,6 +132,7 @@ trait JavaModule extends mill.Module with TaskModule { outer =>
       upstreamCompileOutput()
     )
   }
+
   def localClasspath = T{
     resources() ++ Agg(compile().classes)
   }
@@ -158,19 +161,23 @@ trait JavaModule extends mill.Module with TaskModule { outer =>
     * upstream dependencies do not change
     */
   def upstreamAssembly = T{
-    createAssembly(upstreamAssemblyClasspath().map(_.path), mainClass())
+    createAssembly(
+      upstreamAssemblyClasspath().map(_.path),
+      mainClass(),
+      assemblyRules = assemblyRules
+    )
   }
 
   def assembly = T{
     createAssembly(
       Agg.from(localClasspath().map(_.path)),
       mainClass(),
       prependShellScript(),
-      Some(upstreamAssembly().path)
+      Some(upstreamAssembly().path),
+      assemblyRules
     )
   }
 
-
   def jar = T{
     createJar(
       localClasspath().map(_.path).filter(exists),

scalalib/test/resources/hello-world-multi/core/resources/reference.conf

@@ -0,0 +1,4 @@
+##############################
+# Core Reference Config File #
+##############################
+bar.baz=hello

scalalib/test/resources/hello-world-multi/core/src/Main.scala

@@ -0,0 +1,5 @@
+object Main extends App {
+  val person = Person.fromString("rockjam:25")
+  println(s"hello ${person.name}, your age is: ${person.age}")
+}
+

scalalib/test/resources/hello-world-multi/model/resources/reference.conf

@@ -0,0 +1,4 @@
+###############################
+# Model Reference Config File #
+###############################
+foo.bar=2

scalalib/test/resources/hello-world-multi/model/src/Person.scala

@@ -0,0 +1,8 @@
+object Person {
+  def fromString(s: String): Person = {
+    val Array(name, age) = s.split(":")
+    Person(name, age.toInt)
+  }
+}
+case class Person(name: String, age: Int)
+

scalalib/test/resources/hello-world/core/resources/reference.conf

@@ -0,0 +1,4 @@
+########################################
+# My application Reference Config File #
+########################################
+akka.http.client.user-agent-header="hello-world-client"