From 266616e60e9431cb4c8ade426fb5177d71c229e8 Mon Sep 17 00:00:00 2001
From: John Johnson II <john@potenciasoftware.com>
Date: Mon, 30 May 2022 21:00:18 -0600
Subject: [PATCH] Added the ability to run a custom logic upon quitting the
 REPL

---
 build.sbt                                     |  1 +
 .../com/potenciasoftware/rebel/BaseRepl.scala | 61 ++++++++++++++++++-
 .../potenciasoftware/rebel/BaseReplTest.scala | 40 +++++++++++-
 3 files changed, 100 insertions(+), 2 deletions(-)

diff --git a/build.sbt b/build.sbt
index 11a3cdc..40f08ef 100644
--- a/build.sbt
+++ b/build.sbt
@@ -21,6 +21,7 @@ lazy val rebel = (project in file("."))
   .settings(
     name := "rebel",
     libraryDependencies += "org.scala-lang" % "scala-compiler" % scalaVersion.value,
+    libraryDependencies += "dev.zio" %% "zio" % "2.0.0-RC6",
     libraryDependencies += scalaTest % Test,
   )
 
diff --git a/src/main/scala/com/potenciasoftware/rebel/BaseRepl.scala b/src/main/scala/com/potenciasoftware/rebel/BaseRepl.scala
index 1a601ea..f3d65b5 100644
--- a/src/main/scala/com/potenciasoftware/rebel/BaseRepl.scala
+++ b/src/main/scala/com/potenciasoftware/rebel/BaseRepl.scala
@@ -1,12 +1,16 @@
 package com.potenciasoftware.rebel
 
+import zio._
+
 import java.lang.reflect.Field
 import java.net.URLClassLoader
 import scala.reflect.ClassTag
 import scala.reflect.runtime.universe.TypeTag
 import scala.tools.nsc.Settings
 import scala.tools.nsc.interpreter.Repl
+import scala.tools.nsc.interpreter.shell.Completion
 import scala.tools.nsc.interpreter.shell.ILoop
+import scala.tools.nsc.interpreter.shell.NoCompletion
 import scala.tools.nsc.interpreter.shell.ShellConfig
 import scala.tools.nsc.typechecker.TypeStrings
 
@@ -54,7 +58,7 @@ class BaseRepl {
    * Override to provide bound values.
    * These will be available from within the REPL.
    */
-  protected def boundValues: Seq[Parameter] = Seq()
+  protected def boundValues: Seq[Parameter] = Seq.empty
 
   // Because ILoop declares 'prompt' to be a lazy val,
   // we can't just override it or set it directly.
@@ -76,6 +80,12 @@ class BaseRepl {
    */
   protected def startupScript: String = ""
 
+  /** Override to provide additional colon commands to the REPL. */
+  protected def customCommands: Seq[LoopCommand] = Seq.empty
+
+  /** Override to provide logic to execute when quitting the REPL. */
+  def onQuit(): Unit = ()
+
   /** Read the current text of the REPL prompt. */
   def prompt: String = repl.prompt
 
@@ -103,6 +113,26 @@ class BaseRepl {
           intp.interpret(startupScript)
       }
     }
+
+    private def customQuit(q: LoopCommand): Seq[LoopCommand] =
+      Seq(LoopCommand.cmd(
+        name = q.name,
+        usage = q.usage,
+        help = q.help,
+        f = { line =>
+          delayedAction(5.seconds) { sys.exit() }
+          onQuit()
+          q(line)
+        },
+        completion = q.completion))
+
+    override def commands: List[LoopCommand] = {
+      val (Seq(quitCommand), cmds) =
+        super.commands.partition(_.name == "quit")
+      (cmds ++ customCommands
+        .map(_.convert(LoopCommand.cmd _, Result.apply)))
+        .sortBy(_.name) ++ customQuit(quitCommand)
+    }
   }
 
   def run(): Unit = {
@@ -111,8 +141,37 @@ class BaseRepl {
 }
 
 object BaseRepl {
+
   private val WelcomePlaceholder = "%%%%welcome%%%%"
 
+  private def delayedAction(after: Duration)(action: => Unit): Unit =
+    Runtime.default.unsafeRunAsync {
+      ZIO.attempt(action)
+        .delay(after)
+        .sandbox
+        .catchAll(_ => ZIO.unit)
+    }
+
+  case class LoopCommand(
+    name: String,
+    usage: String,
+    help: String,
+    f: String => LoopCommand.Result,
+    completion: Completion = NoCompletion
+  ) {
+    private[BaseRepl] def convert[A, B](
+      toCommand: (String, String, String, String => B, Completion) => A,
+      toResult: (Boolean, Option[String]) => B
+    ): A =
+      toCommand(name, usage, help,
+      { s => val r = f(s); toResult(r.keepRunning, r.lineToRecord) },
+      completion)
+  }
+
+  object LoopCommand {
+    case class Result(keepRunning: Boolean, lineToRecord: Option[String])
+  }
+
   class Parameter private (
     name: String,
     `type`: String,
diff --git a/src/test/scala/com/potenciasoftware/rebel/BaseReplTest.scala b/src/test/scala/com/potenciasoftware/rebel/BaseReplTest.scala
index 459bcdb..dbe4933 100644
--- a/src/test/scala/com/potenciasoftware/rebel/BaseReplTest.scala
+++ b/src/test/scala/com/potenciasoftware/rebel/BaseReplTest.scala
@@ -8,7 +8,7 @@ import scala.tools.nsc.Settings
 import scala.tools.nsc.interpreter.shell.ILoop
 import scala.tools.nsc.interpreter.shell.ShellConfig
 
-import BaseRepl.Parameter
+import BaseRepl._
 import BaseReplTest._
 import TestUtils._
 
@@ -117,6 +117,44 @@ class BaseReplTest extends AnyFlatSpec with Matchers {
         """
         |scala> printAnswer()The answer is: 42""".stripMargin
   }
+
+  def customCommand() = new TestRepl { repl =>
+    override protected def customCommands: Seq[LoopCommand] =
+      Seq(LoopCommand(
+        "ps1", "<promptText>", "Change the prompt text",
+        { text =>
+          repl.prompt = "\n" + text
+          LoopCommand.Result(true, None)
+        }))
+  }
+
+  it should "allow custom commands" in {
+    replTest[BaseReplTest]("customCommand",
+      ":help ps1", ":ps1 $", "1+1"
+    ).out.map(_.trim).asBlock shouldBe
+    """
+    |scala>
+    |Change the prompt text
+    |
+    |scala>
+    |$val res0: Int = 2
+    |
+    |$""".stripMargin
+  }
+
+  def customQuit() = new TestRepl {
+    override def onQuit(): Unit = {
+      import zio._
+      print("Quitting")
+      // The quit command will only wait 5 seconds before issuing a sys.exit()
+      Thread.sleep(1.minute.toMillis)
+      println("...")
+    }
+  }
+
+  it should "allow custom logic during quit (up to 5 seconds)" in {
+    replTest[BaseReplTest]("customQuit").out(1) shouldBe "scala> Quitting"
+  }
 }
 
 object BaseReplTest {