Parse args directly in MainModule commands to allow for ``--all`-style argument parsing

3 files changed, 70 insertions, 74 deletions

diff --git a/main/src/mill/main/MainModule.scala b/main/src/mill/main/MainModule.scala
index bd8befb3..3ec99f12 100644
--- a/main/src/mill/main/MainModule.scala
+++ b/main/src/mill/main/MainModule.scala
@@ -6,56 +6,70 @@ import pprint.{Renderer, Truncated}
 trait MainModule extends mill.Module{
 
   implicit def millDiscover: mill.define.Discover[_]
-  implicit def millScoptTargetReads[T] = new mill.main.Tasks.Scopt[T]()
+  implicit def millScoptTasksReads[T] = new mill.main.Tasks.Scopt[T]()
   implicit def millScoptEvaluatorReads[T] = new mill.main.EvaluatorScopt[T]()
 
-  def resolve(targets: mill.main.Tasks[Any]*) = mill.T.command{
-    targets.flatMap(_.value).foreach(println)
+  def resolve(evaluator: mill.eval.Evaluator[Any], targets: String*) = mill.T.command{
+    RunScript.resolveTargets(evaluator, targets, multiSelect = true) match{
+      case Left(err) => mill.eval.Result.Failure(err)
+      case Right(tasks) => mill.eval.Result.Success(tasks.foreach(println))
+    }
   }
-  def describe(evaluator: mill.eval.Evaluator[Any],
-               targets: mill.main.Tasks[Any]*) = mill.T.command{
-    for{
-      t <- targets
-      target <- t.value
-      tree = ReplApplyHandler.pprintTask(target, evaluator)
-      val defaults = pprint.PPrinter()
-      val renderer = new Renderer(
-        defaults.defaultWidth,
-        defaults.colorApplyPrefix,
-        defaults.colorLiteral,
-        defaults.defaultIndent
-      )
-      val rendered = renderer.rec(tree, 0, 0).iter
-      val truncated = new Truncated(rendered, defaults.defaultWidth, defaults.defaultHeight)
-      str <- truncated ++ Iterator("\n")
-    } {
-      print(str)
+
+  def describe(evaluator: mill.eval.Evaluator[Any], targets: String*) = mill.T.command{
+    RunScript.resolveTargets(evaluator, targets, multiSelect = true) match{
+      case Left(err) => mill.eval.Result.Failure(err)
+      case Right(tasks) =>
+        for{
+          task <- tasks
+          tree = ReplApplyHandler.pprintTask(task, evaluator)
+          val defaults = pprint.PPrinter()
+          val renderer = new Renderer(
+            defaults.defaultWidth,
+            defaults.colorApplyPrefix,
+            defaults.colorLiteral,
+            defaults.defaultIndent
+          )
+          val rendered = renderer.rec(tree, 0, 0).iter
+          val truncated = new Truncated(rendered, defaults.defaultWidth, defaults.defaultHeight)
+          str <- truncated ++ Iterator("\n")
+        } {
+          print(str)
+        }
+        mill.eval.Result.Success(())
     }
   }
+
   def all(evaluator: mill.eval.Evaluator[Any],
-          targets: mill.main.Tasks[Any]*) = mill.T.command{
-    val (watched, res) = RunScript.evaluate(
-      evaluator,
-      mill.util.Strict.Agg.from(targets.flatMap(_.value))
-    )
-    Watched((), watched)
+          targets: String*) = mill.T.command{
+    RunScript.evaluateTarget(evaluator, targets, multiSelect = true) match{
+      case Left(err) => mill.eval.Result.Failure(err)
+      case Right((watched, Left(err))) => mill.eval.Result.Failure(err, Some(Watched((), watched)))
+      case Right((watched, Right(res))) => mill.eval.Result.Success(Watched((), watched))
+    }
   }
+
   def show(evaluator: mill.eval.Evaluator[Any],
-           targets: mill.main.Tasks[Any]*) = mill.T.command{
-    val (watched, res) = mill.main.RunScript.evaluate(
-      // When using `show`, redirect all stdout of the evaluated tasks so the
-      // printed JSON is the only thing printed to stdout.
+           targets: String*) = mill.T.command{
+    RunScript.evaluateTarget(
       evaluator.copy(
+        // When using `show`, redirect all stdout of the evaluated tasks so the
+        // printed JSON is the only thing printed to stdout.
         log = evaluator.log match{
           case PrintLogger(c1, c2, o, i, e) => PrintLogger(c1, c2, e, i, e)
           case l => l
         }
       ),
-      mill.util.Strict.Agg.from(targets.flatMap(_.value))
-    )
-    for(json <- res.right.get.flatMap(_._2)){
-      println(json)
+      targets,
+      multiSelect = false
+    ) match{
+      case Left(err) => mill.eval.Result.Failure(err)
+      case Right((watched, Left(err))) => mill.eval.Result.Failure(err, Some(Watched((), watched)))
+      case Right((watched, Right(res))) =>
+        for(json <- res.flatMap(_._2)){
+          println(json)
+        }
+        mill.eval.Result.Success(Watched((), watched))
     }
-    Watched((), watched)
   }
 }
diff --git a/main/src/mill/main/MainScopts.scala b/main/src/mill/main/MainScopts.scala
index acf48d05..f8f99c88 100644
--- a/main/src/mill/main/MainScopts.scala
+++ b/main/src/mill/main/MainScopts.scala
@@ -1,43 +1,20 @@
 package mill.main
-import mill.define.ExternalModule
-import mill.eval.{Evaluator, PathRef}
-import mill.util.ParseArgs
+import mill.eval.{Evaluator}
 
 case class Tasks[T](value: Seq[mill.define.NamedTask[T]])
 
 object Tasks{
 
-  class Scopt[T]()
-    extends scopt.Read[Tasks[T]] {
-    def arity = 0
+  class Scopt[T]() extends scopt.Read[Tasks[T]] {
+    def arity = 1
 
     def reads = s => {
-      val rootModule = Evaluator.currentEvaluator.get.rootModule
-      val d = rootModule.millDiscover
-      val (expanded, leftover) = ParseArgs(Seq(s)).fold(e => throw new Exception(e), identity)
-      val resolved = expanded.map {
-        case (Some(scoping), segments) =>
-          val moduleCls = rootModule.getClass.getClassLoader.loadClass(scoping.render + "$")
-          val externalRootModule = moduleCls.getField("MODULE$").get(moduleCls).asInstanceOf[ExternalModule]
-          val crossSelectors = segments.value.map {
-            case mill.define.Segment.Cross(x) => x.toList.map(_.toString)
-            case _ => Nil
-          }
-          mill.main.Resolve.resolve(segments.value.toList, externalRootModule, d, leftover, crossSelectors.toList, Nil)
-        case (None, segments) =>
-          val crossSelectors = segments.value.map {
-            case mill.define.Segment.Cross(x) => x.toList.map(_.toString)
-            case _ => Nil
-          }
-          mill.main.Resolve.resolve(segments.value.toList, rootModule, d, leftover, crossSelectors.toList, Nil)
-      }
-      mill.util.EitherOps.sequence(resolved) match {
-        case Left(s) => throw new Exception(s)
-        case Right(ts) => Tasks(ts.flatten).asInstanceOf[Tasks[T]]
+      RunScript.resolveTargets(Evaluator.currentEvaluator.get, Seq(s), multiSelect = false) match{
+        case Left(err) => throw new Exception(err)
+        case Right(tasks) => Tasks(tasks).asInstanceOf[Tasks[T]]
       }
     }
   }
-
 }
 
 class EvaluatorScopt[T]()
diff --git a/main/src/mill/main/RunScript.scala b/main/src/mill/main/RunScript.scala
index bd21a821..5b5eb429 100644
--- a/main/src/mill/main/RunScript.scala
+++ b/main/src/mill/main/RunScript.scala
@@ -52,7 +52,7 @@ object RunScript{
 
     val evaluated = for{
       evaluator <- evalRes
-      (evalWatches, res) <- Res(evaluateTarget(evaluator, scriptArgs))
+      (evalWatches, res) <- Res(evaluateTarget(evaluator, scriptArgs, multiSelect = false))
     } yield {
       val alreadyStale = evalWatches.exists(p => p.sig != new PathRef(p.path, p.quick).sig)
       // If the file changed between the creation of the original
@@ -129,11 +129,13 @@ object RunScript{
     } yield (module, discover)
   }
 
-  def evaluateTarget[T](evaluator: Evaluator[T], scriptArgs: Seq[String]) = {
+  def resolveTargets[T](evaluator: Evaluator[T],
+                        scriptArgs: Seq[String],
+                        multiSelect: Boolean) = {
     for {
-      parsed <- ParseArgs(scriptArgs)
+      parsed <- ParseArgs(scriptArgs, multiSelect = multiSelect)
       (selectors, args) = parsed
-      targets <- {
+      taskss <- {
         val selected = selectors.map { case (scopedSel, sel) =>
           val (rootModule, discover) = scopedSel match{
             case None => (evaluator.rootModule, evaluator.discover)
@@ -167,11 +169,14 @@ object RunScript{
         }
         EitherOps.sequence(selected)
       }
-    } yield {
-      val (watched, res) = evaluate(
-        evaluator,
-        Agg.from(targets.flatten.distinct)
-      )
+    } yield taskss.flatten
+  }
+
+  def evaluateTarget[T](evaluator: Evaluator[T],
+                        scriptArgs: Seq[String],
+                        multiSelect: Boolean) = {
+    for (targets <- resolveTargets(evaluator, scriptArgs, multiSelect)) yield {
+      val (watched, res) = evaluate(evaluator, Agg.from(targets.distinct))
 
       val watched2 = for{
         x <- res.right.toSeq