Allow for executions to be passed through and throw exception on help

twitter · Feb 3, 2016 · 507e06a · 507e06a
1 parent d4ef80b
commit 507e06a
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 14 deletions.
diff --git a/scalding-core/src/main/scala/com/twitter/scalding/ArgHelp.scala b/scalding-core/src/main/scala/com/twitter/scalding/ArgHelp.scala
@@ -10,25 +10,27 @@ case class OptionalArg(key: String, description: String) extends DescribedArg
 case class ListArg(key: String, description: String) extends DescribedArg
 case class BooleanArg(key: String, description: String) extends DescribedArg
 
+class HelpException extends RuntimeException("User asked for help")
+class DescriptionValidationException(msg: String) extends RuntimeException(msg)
+
 trait ArgHelper {
   /**
    * Similar to describe but validate all args are described
    * @param describedArgs List of Argument Descriptions
    * @param ex Input Execution
    * @return Output Execution
    */
-  def validatedDescribe(describedArgs: Seq[DescribedArg], ex: Execution[Unit]): Execution[Unit] = {
+  def validatedDescribe[T](describedArgs: Seq[DescribedArg], ex: Execution[T]): Execution[T] = {
     val describedKeys = describedArgs.map(_.key).toSet
     Execution.getArgs.flatMap { args =>
-      val describedEx = describe(describedArgs, ex)
       val missingKeys = args.m.keySet.filter(_.nonEmpty).diff(describedKeys)
 
       if (args.boolean("help")) {
         helpRequest(describedArgs)
       } else {
         if (missingKeys.nonEmpty) {
           val msg = missingKeys.mkString(", ")
-          Execution.failed(throw new RuntimeException(s"Must describe missing keys : $msg"))
+          Execution.failed(throw new DescriptionValidationException(s"Must describe missing keys : $msg"))
         } else {
           ex
         }
@@ -44,7 +46,7 @@ trait ArgHelper {
    * @param ex Input Execution
    * @return Output Execution
    */
-  def describe(describedArgs: Seq[DescribedArg], ex: Execution[Unit]): Execution[Unit] = {
+  def describe[T](describedArgs: Seq[DescribedArg], ex: Execution[T]): Execution[T] = {
     Execution.getArgs.flatMap { args =>
       if (args.boolean("help")) {
         helpRequest(describedArgs)
@@ -54,14 +56,14 @@ trait ArgHelper {
     }
   }
 
-  def helpRequest(describedArgs: Seq[DescribedArg]): Execution[Unit] = {
+  def helpRequest[T](describedArgs: Seq[DescribedArg]): Execution[T] = {
     val top = "\n###########################################################################\n\n"
     val usage = s"Command Line Args :: ${argString(describedArgs)}\n\n\n"
     val bottom = "\n\n###########################################################################\n"
 
     println(top + usage + help(describedArgs) + bottom)
 
-    Execution.unit
+    Execution.failed(throw new HelpException())
   }
 
   /**

diff --git a/scalding-core/src/test/scala/com/twitter/scalding/ArgHelpTest.scala b/scalding-core/src/test/scala/com/twitter/scalding/ArgHelpTest.scala
@@ -3,14 +3,14 @@ package com.twitter.scalding
 import org.scalatest.{ Matchers, WordSpec }
 
 case class ArgHelperTest(testFn: Seq[DescribedArg] => Unit) extends ArgHelper {
-  override def helpRequest(describedArgs: Seq[DescribedArg]): Execution[Unit] = {
+  override def helpRequest[T](describedArgs: Seq[DescribedArg]): Execution[T] = {
     testFn(describedArgs)
-    Execution.unit
+    Execution.failed(throw new HelpException())
   }
 }
 
 class ArgHelpTest extends WordSpec with Matchers {
-  def job = TypedPipe.from(List(1, 2, 3)).toIterableExecution.map(_.foreach(println))
+  def job = TypedPipe.from(List(1, 2, 3)).toIterableExecution
 
   "ArgHelper" should {
     "print help when asked" in {
@@ -20,7 +20,9 @@ class ArgHelpTest extends WordSpec with Matchers {
       val args = List(OptionalArg("name", "Name of person"))
       val config = Config.unitTestDefault.setArgs(Args("--help"))
 
-      helper.describe(args, job.unit).waitFor(config, Local(true))
+      intercept[HelpException] {
+        helper.describe(args, job).waitFor(config, Local(true)).get
+      }
       assert(helpCalled, "Help function was called")
     }
   }
@@ -33,8 +35,9 @@ class ArgHelpTest extends WordSpec with Matchers {
       val args = List(OptionalArg("name", "Name of person"))
       val config = Config.unitTestDefault.setArgs(Args(""))
 
-      helper.describe(args, job.unit).waitFor(config, Local(true))
+      val returnValues = helper.describe(args, job).waitFor(config, Local(true)).get.toList
       assert(!helpCalled, "Help function was not called")
+      assert(returnValues == List(1, 2, 3))
     }
   }
 
@@ -46,7 +49,9 @@ class ArgHelpTest extends WordSpec with Matchers {
       val args = List(OptionalArg("name", "Name of person"))
       val config = Config.unitTestDefault.setArgs(Args(List("--help", "--name", "Bill", "--phone", "111")))
 
-      helper.validatedDescribe(args, job.unit).waitFor(config, Local(true))
+      intercept[HelpException] {
+        helper.validatedDescribe(args, job).waitFor(config, Local(true)).get
+      }
       assert(helpCalled, "Help was output")
     }
   }
@@ -56,7 +61,8 @@ class ArgHelpTest extends WordSpec with Matchers {
       val args = List(OptionalArg("name", "Name of person"), OptionalArg("phone", "Person's phone"))
       val config = Config.unitTestDefault.setArgs(Args(List("--name", "Bill", "--phone", "111")))
 
-      assert(ArgHelp.validatedDescribe(args, job.unit).waitFor(config, Local(true)).isSuccess)
+      val returnValues = ArgHelp.validatedDescribe(args, job).waitFor(config, Local(true)).get
+      assert(returnValues == List(1, 2, 3))
     }
   }
 
@@ -65,7 +71,9 @@ class ArgHelpTest extends WordSpec with Matchers {
       val args = List(OptionalArg("name", "Name of person"), OptionalArg("phone", "Person's phone"))
       val config = Config.unitTestDefault.setArgs(Args(List("--name", "Bill", "--phone", "111", "--address", "123")))
 
-      assert(ArgHelp.validatedDescribe(args, job.unit).waitFor(config, Local(true)).isFailure)
+      intercept[DescriptionValidationException] {
+        ArgHelp.validatedDescribe(args, job.unit).waitFor(config, Local(true)).get
+      }
     }
   }
 }