fix @main method handling in scripts run via --class-based

com-lihaoyi · Dec 19, 2019 · b6ae3ff · b6ae3ff
1 parent 0aad8e0
commit b6ae3ff
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 9 deletions.
diff --git a/amm/interp/src/main/scala/ammonite/interp/Interpreter.scala b/amm/interp/src/main/scala/ammonite/interp/Interpreter.scala
@@ -43,7 +43,7 @@ class Interpreter(val printer: Printer,
                   val createFrame: () => Frame,
                   initialClassLoader: ClassLoader = null,
                   replCodeWrapper: CodeWrapper,
-                  scriptCodeWrapper: CodeWrapper,
+                  val scriptCodeWrapper: CodeWrapper,
                   alreadyLoadedDependencies: Seq[Dependency],
                   importHooks: Map[Seq[String], ImportHook] = ImportHook.defaults,
                   classPathWhitelist: Set[Seq[String]] = Set.empty)

diff --git a/amm/src/main/scala/ammonite/main/Scripts.scala b/amm/src/main/scala/ammonite/main/Scripts.scala
@@ -74,20 +74,35 @@ object Scripts {
           .evalClassloader
           .loadClass(processed.blockInfo.last.id.wrapperPath + "$")
 
-      routesCls =
-        interp
-          .evalClassloader
-          .loadClass(routeClsName + "$$routes$")
-
-      scriptMains =
-        routesCls
+      scriptMains = interp.scriptCodeWrapper match{
+        case ammonite.interp.CodeWrapper =>
+          interp
+            .evalClassloader
+            .loadClass(routeClsName + "$$routes$")
             .getField("MODULE$")
             .get(null)
             .asInstanceOf[() => Seq[Router.EntryPoint[Any]]]
             .apply()
 
+        case ammonite.interp.CodeClassWrapper =>
+          val outer = interp
+            .evalClassloader
+            .loadClass(routeClsName)
+            .getMethod("instance")
+            .invoke(null)
+
+          outer.getClass.getMethod("$routes").invoke(outer)
+            .asInstanceOf[() => Seq[Router.EntryPoint[Any]]]
+            .apply()
+        case _ => Nil
+      }
 
-      mainObj = mainCls.getField("MODULE$").get(null)
+      mainObj = interp.scriptCodeWrapper match {
+        case ammonite.interp.CodeWrapper => mainCls.getField("MODULE$").get(null)
+        case ammonite.interp.CodeClassWrapper =>
+          mainCls.getMethod("instance").invoke(mainCls.getField("MODULE$").get(null))
+        case _ => ()
+      }
 
       res <- Util.withContextClassloader(interp.evalClassloader){
         scriptMains match {

diff --git a/readme/Footer.scalatex b/readme/Footer.scalatex
@@ -123,6 +123,10 @@
 
 
   @sect{Changelog}
+    @sect{1.9.1}
+      @ul
+        @li
+          Fix @code{@@main} method handling in scripts run via @code{--class-based}
     @sect{1.9.0}
       @ul
         @li