DRY up task resolution between MainModule and RunScript

3 files changed, 60 insertions, 58 deletions

diff --git a/main/src/mill/main/MainModule.scala b/main/src/mill/main/MainModule.scala
index 8a4cbefa..b1068c29 100644
--- a/main/src/mill/main/MainModule.scala
+++ b/main/src/mill/main/MainModule.scala
@@ -8,7 +8,7 @@ import upickle.Js
 object MainModule{
   def resolveTasks[T](evaluator: Evaluator[Any], targets: Seq[String], multiSelect: Boolean)
                      (f: List[NamedTask[Any]] => T) = {
-    RunScript.resolveTasks(evaluator, targets, multiSelect) match{
+    RunScript.resolveTasks(mill.main.Resolve, evaluator, targets, multiSelect) match{
       case Left(err) => Result.Failure(err)
       case Right(tasks) => Result.Success(f(tasks))
     }
@@ -32,30 +32,12 @@ trait MainModule extends mill.Module{
   implicit def millScoptEvaluatorReads[T] = new mill.main.EvaluatorScopt[T]()
 
   def resolve(evaluator: Evaluator[Any], targets: String*) = mill.T.command{
-    val resolved = for {
-      parsed <- ParseArgs(targets, multiSelect = true)
-      (selectors, args) = parsed
-      taskss <- {
-        val selected = selectors.map { case (scopedSel, sel) =>
-          val (rootModule, discover, crossSelectors) =
-            mill.main.RunScript.prepareResolve(evaluator, scopedSel, sel)
-
-          try {
-            mill.eval.Evaluator.currentEvaluator.set(evaluator)
-            mill.main.ResolveMetadata.resolve(
-              sel.value.toList, rootModule, discover,
-              args, crossSelectors.toList, Nil
-            )
-          } finally{
-            mill.eval.Evaluator.currentEvaluator.set(null)
-          }
-        }
-        EitherOps.sequence(selected)
-      }
-    } yield taskss.flatten
+    val resolved = RunScript.resolveTasks(
+      mill.main.ResolveMetadata, evaluator, targets, multiSelect = true
+    )
     resolved match{
-      case Left(err) => throw new Exception(err)
-      case Right(r) => r.foreach(println)
+      case Left(err) => Result.Failure(err)
+      case Right(r) => Result.Success(r.foreach(println))
     }
   }
 
diff --git a/main/src/mill/main/MainScopts.scala b/main/src/mill/main/MainScopts.scala
index 653a299d..36e13717 100644
--- a/main/src/mill/main/MainScopts.scala
+++ b/main/src/mill/main/MainScopts.scala
@@ -1,5 +1,5 @@
 package mill.main
-import mill.eval.{Evaluator}
+import mill.eval.Evaluator
 
 case class Tasks[T](value: Seq[mill.define.NamedTask[T]])
 
@@ -9,7 +9,12 @@ object Tasks{
     def arity = 1
 
     def reads = s => {
-      RunScript.resolveTasks(Evaluator.currentEvaluator.get, Seq(s), multiSelect = false) match{
+      RunScript.resolveTasks(
+        mill.main.Resolve,
+        Evaluator.currentEvaluator.get,
+        Seq(s),
+        multiSelect = false
+      ) match{
         case Left(err) => throw new Exception(err)
         case Right(tasks) => Tasks(tasks).asInstanceOf[Tasks[T]]
       }
diff --git a/main/src/mill/main/RunScript.scala b/main/src/mill/main/RunScript.scala
index c9155551..c00b6a14 100644
--- a/main/src/mill/main/RunScript.scala
+++ b/main/src/mill/main/RunScript.scala
@@ -14,6 +14,8 @@ import mill.util.{EitherOps, Logger, ParseArgs, Watched}
 import mill.util.Strict.Agg
 import upickle.Js
 
+import scala.reflect.ClassTag
+
 /**
   * Custom version of ammonite.main.Scripts, letting us run the build.sc script
   * directly without going through Ammonite's main-method/argument-parsing
@@ -129,58 +131,71 @@ object RunScript{
     } yield (module, discover)
   }
 
-  def resolveTasks[T](evaluator: Evaluator[T],
-                      scriptArgs: Seq[String],
-                      multiSelect: Boolean) = {
+  def resolveTasks[T, R: ClassTag](resolver: mill.main.Resolve[R],
+                                   evaluator: Evaluator[T],
+                                   scriptArgs: Seq[String],
+                                   multiSelect: Boolean) = {
     for {
       parsed <- ParseArgs(scriptArgs, multiSelect = multiSelect)
       (selectors, args) = parsed
       taskss <- {
         val selected = selectors.map { case (scopedSel, sel) =>
-          val (rootModule, discover, crossSelectors) =
-            prepareResolve(evaluator, scopedSel, sel)
-
-          try {
-            // We inject the `evaluator.rootModule` into the TargetScopt, rather
-            // than the `rootModule`, because even if you are running an external
-            // module we still want you to be able to resolve targets from your
-            // main build. Resolving targets from external builds as CLI arguments
-            // is not currently supported
-            mill.eval.Evaluator.currentEvaluator.set(evaluator)
-            mill.main.Resolve.resolve(
-              sel.value.toList, rootModule, discover,
-              args, crossSelectors.toList, Nil
-            )
-          } finally{
-            mill.eval.Evaluator.currentEvaluator.set(null)
+          for(res <- prepareResolve(evaluator, scopedSel, sel))
+          yield {
+            val (rootModule, discover, crossSelectors) = res
+
+
+            try {
+              // We inject the `evaluator.rootModule` into the TargetScopt, rather
+              // than the `rootModule`, because even if you are running an external
+              // module we still want you to be able to resolve targets from your
+              // main build. Resolving targets from external builds as CLI arguments
+              // is not currently supported
+              mill.eval.Evaluator.currentEvaluator.set(evaluator)
+              resolver.resolve(
+                sel.value.toList, rootModule, discover,
+                args, crossSelectors.toList, Nil
+              )
+            } finally {
+              mill.eval.Evaluator.currentEvaluator.set(null)
+            }
           }
         }
         EitherOps.sequence(selected)
       }
-    } yield taskss.flatten
+      res <- EitherOps.sequence(taskss)
+    } yield res.flatten
   }
 
+
   def prepareResolve[T](evaluator: Evaluator[T], scopedSel: Option[Segments], sel: Segments) = {
-    val (rootModule, discover) = scopedSel match {
-      case None => (evaluator.rootModule, evaluator.discover)
-      case Some(scoping) =>
+    for {
+      res <- scopedSel match {
+        case None => Right((evaluator.rootModule, evaluator.discover))
+        case Some (scoping) =>
         val moduleCls =
-          evaluator.rootModule.getClass.getClassLoader.loadClass(scoping.render + "$")
+          evaluator.rootModule.getClass.getClassLoader.loadClass (scoping.render + "$")
 
-        val rootModule = moduleCls.getField("MODULE$").get(moduleCls).asInstanceOf[ExternalModule]
-        (rootModule, rootModule.millDiscover)
-    }
-    val crossSelectors = sel.value.map {
-      case Segment.Cross(x) => x.toList.map(_.toString)
-      case _ => Nil
+        moduleCls.getField("MODULE$").get(moduleCls) match {
+          case rootModule: ExternalModule => Right ((rootModule, rootModule.millDiscover))
+          case _ => Left("External Module " + scoping.render + " not found")
+        }
+
+      }
+    } yield {
+      val (rootModule, discover) = res
+      val crossSelectors = sel.value.map {
+        case Segment.Cross(x) => x.toList.map(_.toString)
+        case _ => Nil
+      }
+      (rootModule, discover, crossSelectors)
     }
-    (rootModule, discover, crossSelectors)
   }
 
   def evaluateTasks[T](evaluator: Evaluator[T],
                        scriptArgs: Seq[String],
                        multiSelect: Boolean) = {
-    for (targets <- resolveTasks(evaluator, scriptArgs, multiSelect)) yield {
+    for (targets <- resolveTasks(mill.main.Resolve, evaluator, scriptArgs, multiSelect)) yield {
       val (watched, res) = evaluate(evaluator, Agg.from(targets.distinct))
 
       val watched2 = for{