Partial support for ParserForClass

com-lihaoyi · Feb 28, 2022 · 3b33efc · 3b33efc
1 parent 58b6f17
commit 3b33efc
Show file tree

Hide file tree

Showing 17 changed files with 405 additions and 267 deletions.
diff --git a/mainargs/src-2/ParserForClassCompanionVersionSpecific.scala b/mainargs/src-2/ParserForClassCompanionVersionSpecific.scala
@@ -4,6 +4,6 @@ import acyclic.skipped
 
 import scala.language.experimental.macros
 
-private [mainargs] trait ParserForClassCompanionVersionSpecific {
+private[mainargs] trait ParserForClassCompanionVersionSpecific {
   def apply[T]: ParserForClass[T] = macro Macros.parserForClass[T]
 }
diff --git a/mainargs/src-2/ParserForMethodsCompanionVersionSpecific.scala b/mainargs/src-2/ParserForMethodsCompanionVersionSpecific.scala
@@ -4,6 +4,6 @@ import acyclic.skipped
 
 import scala.language.experimental.macros
 
-private [mainargs] trait ParserForMethodsCompanionVersionSpecific {
+private[mainargs] trait ParserForMethodsCompanionVersionSpecific {
   def apply[B](base: B): ParserForMethods[B] = macro Macros.parserForMethods[B]
 }
diff --git a/mainargs/src-3/Macros.scala b/mainargs/src-3/Macros.scala
@@ -3,48 +3,17 @@ package mainargs
 import scala.quoted._
 
 object Macros {
+  private def mainAnnotation(using Quotes) = quotes.reflect.TypeRepr.of[mainargs.main].typeSymbol
+  private def argAnnotation(using Quotes) = quotes.reflect.TypeRepr.of[mainargs.arg].typeSymbol
   def parserForMethods[B](base: Expr[B])(using Quotes, Type[B]): Expr[ParserForMethods[B]] = {
     import quotes.reflect._
     val allMethods = TypeRepr.of[B].typeSymbol.memberMethods
-    val mainAnnotation = TypeRepr.of[mainargs.main].typeSymbol
-    val argAnnotation = TypeRepr.of[mainargs.arg].typeSymbol
     val annotatedMethodsWithMainAnnotations = allMethods.flatMap { methodSymbol =>
       methodSymbol.getAnnotation(mainAnnotation).map(methodSymbol -> _)
     }.sortBy(_._1.pos.map(_.start))
-    val mainDatasExprs: Seq[Expr[MainData[Any, B]]] = annotatedMethodsWithMainAnnotations.map { (annotatedMethod, mainAnnotation) =>
-      val params = annotatedMethod.paramSymss.headOption.getOrElse(throw new Exception("Multiple parameter lists not supported"))
-      val defaultParams = getDefaultParams(annotatedMethod)
-      val argSigs = Expr.ofList(params.map { param =>
-        val paramTree = param.tree.asInstanceOf[ValDef]
-        val paramTpe = paramTree.tpt.tpe
-        val arg = param.getAnnotation(argAnnotation).map(_.asExpr.asInstanceOf[Expr[mainargs.arg]]).getOrElse('{ new mainargs.arg() })
-        val paramType = paramTpe.asType
-        paramType match
-          case '[t] =>
-            val defaultParam: Expr[Option[B => t]] = defaultParams.get(param) match {
-              case Some(v) => '{ Some(((_: B) => $v).asInstanceOf[B => t]) }
-              case None => '{ None }
-            }
-            val argReader = Expr.summon[mainargs.ArgReader[t]].getOrElse{
-              report.error(
-                s"No mainargs.ArgReader of ${paramTpe.typeSymbol.fullName} found for parameter ${param.name}",
-                param.pos.get
-              )
-              '{ ??? }
-            }
-            '{ ArgSig.create[t, B](${ Expr(param.name) }, ${ arg }, ${ defaultParam })(using ${ argReader }).asInstanceOf[mainargs.ArgSig[Any, B]] }
-      })
-
-      val invokeRaw: Expr[(B, Seq[Any]) => Any] = {
-        def callOf(args: Expr[Seq[Any]]) = call(annotatedMethod, '{ Seq( ${ args }) })
-        '{ (b: B, params: Seq[Any]) =>
-          ${ callOf('{ params }) }
-        }
-      }
-
-      '{ MainData.create[Any, B](${ Expr(annotatedMethod.name) }, ${ mainAnnotation.asExprOf[mainargs.main] }, ${ argSigs }, ${ invokeRaw }) }
-    }
-    val mainDatas = Expr.ofList(mainDatasExprs)
+    val mainDatas = Expr.ofList(annotatedMethodsWithMainAnnotations.map { (annotatedMethod, mainAnnotationInstance) =>
+      createMainData[Any, B](annotatedMethod, mainAnnotationInstance)
+    })
 
     '{
       new ParserForMethods[B](
@@ -53,6 +22,65 @@ object Macros {
     }
   }
 
+  def parserForClass[B](using Quotes, Type[B]): Expr[ParserForClass[B]] = {
+    import quotes.reflect._
+    val typeReprOfB = TypeRepr.of[B]
+    val companionModule = typeReprOfB match {
+      case TypeRef(a,b) => TermRef(a,b)
+    }
+    val typeSymbolOfB = typeReprOfB.typeSymbol
+    val companionModuleType = typeSymbolOfB.companionModule.tree.asInstanceOf[ValDef].tpt.tpe.asType
+    val companionModuleExpr = Ident(companionModule).asExpr
+    val mainAnnotationInstance = typeSymbolOfB.getAnnotation(mainAnnotation).getOrElse {
+      report.error(
+        s"cannot find @main annotation on ${companionModule.name}",
+        typeSymbolOfB.pos.get
+      )
+      ???
+    }
+    val annotatedMethod = TypeRepr.of[B].typeSymbol.companionModule.memberMethod("apply").head
+    companionModuleType match
+      case '[bCompanion] =>
+        val mainData = createMainData[B, bCompanion](annotatedMethod, mainAnnotationInstance)
+        '{
+          new ParserForClass[B](
+            ClassMains[B](${ mainData }.asInstanceOf[MainData[B, Any]], () => ${ Ident(companionModule).asExpr })
+          )
+        }
+  }
+
+  def createMainData[T: Type, B: Type](using Quotes)(method: quotes.reflect.Symbol, annotation: quotes.reflect.Term): Expr[MainData[T, B]] = {
+    import quotes.reflect.*
+    val params = method.paramSymss.headOption.getOrElse(throw new Exception("Multiple parameter lists not supported"))
+    val defaultParams = getDefaultParams(method)
+    val argSigs = Expr.ofList(params.map { param =>
+      val paramTree = param.tree.asInstanceOf[ValDef]
+      val paramTpe = paramTree.tpt.tpe
+      val arg = param.getAnnotation(argAnnotation).map(_.asExpr.asInstanceOf[Expr[mainargs.arg]]).getOrElse('{ new mainargs.arg() })
+      val paramType = paramTpe.asType
+      paramType match
+        case '[t] =>
+          val defaultParam: Expr[Option[B => t]] = defaultParams.get(param) match {
+            case Some(v) => '{ Some(((_: B) => $v).asInstanceOf[B => t]) }
+            case None => '{ None }
+          }
+          val argReader = Expr.summon[mainargs.ArgReader[t]].getOrElse{
+            report.error(
+              s"No mainargs.ArgReader of ###companionModule### found for parameter ${param.name}",
+              param.pos.get
+            )
+            '{ ??? }
+          }
+          '{ (ArgSig.create[t, B](${ Expr(param.name) }, ${ arg }, ${ defaultParam })(using ${ argReader })).asInstanceOf[ArgSig[Any, B]] }
+    })
+
+    val invokeRaw: Expr[(B, Seq[Any]) => T] = {
+      def callOf(args: Expr[Seq[Any]]) = call(method, '{ Seq( ${ args }) })
+      '{ ((b: B, params: Seq[Any]) => ${ callOf('{ params }) }).asInstanceOf[(B, Seq[Any]) => T] }
+    }
+    '{ MainData.create[T, B](${ Expr(method.name) }, ${ annotation.asExprOf[mainargs.main] }, ${ argSigs }, ${ invokeRaw }) }
+  }
+
   /** Call a method given by its symbol.
     *
     * E.g.

diff --git a/mainargs/src-3/ParserForClassCompanionVersionSpecific.scala b/mainargs/src-3/ParserForClassCompanionVersionSpecific.scala
@@ -3,5 +3,5 @@ package mainargs
 import scala.language.experimental.macros
 
 private [mainargs] trait ParserForClassCompanionVersionSpecific {
-  inline def apply[T]: ParserForClass[T] = ???
+  inline def apply[T]: ParserForClass[T] = ${ Macros.parserForClass[T] }
 }
diff --git a/mainargs/test/src-2/ClassTests.scala b/mainargs/test/src-2/ClassTests.scala
diff --git a/mainargs/test/src-2/OldVarargsTests.scala b/mainargs/test/src-2/OldVarargsTests.scala
@@ -1,14 +1,15 @@
 package mainargs
 import utest._
 
-object OldVarargsTests extends VarargsTests{
-  object Base{
+object OldVarargsTests extends VarargsTests {
+  object Base {
 
     @main
     def pureVariadic(nums: Int*) = nums.sum
 
     @main
-    def mixedVariadic(@arg(short = 'f') first: Int, args: String*) = first + args.mkString
+    def mixedVariadic(@arg(short = 'f') first: Int, args: String*) =
+      first + args.mkString
   }
 
   val check = new Checker(ParserForMethods(Base), allowPositional = true)

diff --git a/mainargs/test/src-2/ParserTests.scala b/mainargs/test/src-2/ParserTests.scala
diff --git a/mainargs/test/src-2/VersionSpecific.scala b/mainargs/test/src-2/VersionSpecific.scala
@@ -0,0 +1,5 @@
+package mainargs
+
+object VersionSpecific {
+  val isScala3 = false
+}
diff --git a/mainargs/test/src-3/VersionSpecific.scala b/mainargs/test/src-3/VersionSpecific.scala
@@ -0,0 +1,5 @@
+package mainargs
+
+object VersionSpecific {
+  val isScala3 = true
+}
diff --git a/mainargs/test/src-jvm-2/AmmoniteTests.scala b/mainargs/test/src-jvm-2/AmmoniteTests.scala
@@ -69,7 +69,7 @@ object AmmoniteConfig{
     @arg(doc = "Print this message")
     help: Flag
   )
-  implicit val coreParser: ParserForClass[Core] = ParserForClass[Core]
+  implicit val coreParser = ParserForClass[Core]
 
   @main
   case class Predef(
@@ -86,7 +86,7 @@ object AmmoniteConfig{
         "choose an additional predef to use using `--predef")
     noHomePredef: Flag
   )
-  implicit val predefParser: ParserForClass[Predef] = ParserForClass[Predef]
+  implicit val predefParser = ParserForClass[Predef]
 
   @main
   case class Repl(
@@ -105,12 +105,12 @@ object AmmoniteConfig{
         "friendliness.")
     classBased: Flag
   )
-  implicit val replParser: ParserForClass[Repl] = ParserForClass[Repl]
+  implicit val replParser = ParserForClass[Repl]
 }
 
 
 object AmmoniteTests extends TestSuite{
-  val parser: ParserForClass[AmmoniteConfig] = ParserForClass[AmmoniteConfig]
+  val parser = ParserForClass[AmmoniteConfig]
   val tests = Tests {
 
 

diff --git a/mainargs/test/src-jvm-2/MillTests.scala b/mainargs/test/src-jvm-2/MillTests.scala
@@ -1,7 +1,6 @@
 // package mainargs
 // import utest._
 
-
 // object MillTests extends TestSuite{
 
 //  implicit object PathRead extends TokensReader[os.Path]("path", strs => Right(os.Path(strs.head, os.pwd)))
@@ -107,4 +106,3 @@
 //    }
 //  }
 // }
-