Add new types to enforce wrapping scripts before building project

VirtusLab · May 7, 2023 · caf27c1 · caf27c1
1 parent 5d3e440
commit caf27c1
Show file tree

Hide file tree

Showing 9 changed files with 251 additions and 117 deletions.
diff --git a/modules/build/src/main/scala/scala/build/CrossSources.scala b/modules/build/src/main/scala/scala/build/CrossSources.scala
@@ -28,26 +28,94 @@ import scala.build.testrunner.DynamicTestRunner.globPattern
 import scala.util.Try
 import scala.util.chaining.*
 
-final case class CrossSources(
+/** CrossSources with unwrapped scripts, use [[withWrappedScripts]] to wrap them and obtain an
+  * instance of CrossSources
+  *
+  * See [[CrossSources]] for more information
+  *
+  * @param paths
+  *   paths and realtive paths to sources on disk, wrapped in their build requirements
+  * @param inMemory
+  *   in memory sources (e.g. snippets) wrapped in their build requirements
+  * @param defaultMainClass
+  * @param resourceDirs
+  * @param buildOptions
+  *   build options from sources
+  * @param unwrappedScripts
+  *   in memory script sources, their code must be wrapped before compiling
+  */
+sealed class UnwrappedCrossSources(
   paths: Seq[WithBuildRequirements[(os.Path, os.RelPath)]],
   inMemory: Seq[WithBuildRequirements[Sources.InMemory]],
   defaultMainClass: Option[String],
   resourceDirs: Seq[WithBuildRequirements[os.Path]],
-  buildOptions: Seq[WithBuildRequirements[BuildOptions]]
+  buildOptions: Seq[WithBuildRequirements[BuildOptions]],
+  unwrappedScripts: Seq[WithBuildRequirements[Sources.UnwrappedScript]]
 ) {
 
+  /** For all unwrapped script sources contained in this object wrap them according to provided
+    * BuildOptions
+    *
+    * @param buildOptions
+    *   options used to choose the script wrapper
+    * @return
+    *   CrossSources with all the scripts wrapped
+    */
+  def withWrappedScripts(buildOptions: BuildOptions): CrossSources = {
+    val codeWrapper = ScriptPreprocessor.getScriptWrapper(buildOptions)
+
+    val wrappedScripts = unwrappedScripts.map { unwrapppedWithRequirements =>
+      unwrapppedWithRequirements.map(_.wrap(codeWrapper))
+    }
+
+    CrossSources(
+      paths,
+      inMemory ++ wrappedScripts,
+      defaultMainClass,
+      resourceDirs,
+      this.buildOptions
+    )
+  }
+
   def sharedOptions(baseOptions: BuildOptions): BuildOptions =
     buildOptions
       .filter(_.requirements.isEmpty)
       .map(_.value)
       .foldLeft(baseOptions)(_ orElse _)
 
-  private def needsScalaVersion =
+  protected def needsScalaVersion =
     paths.exists(_.needsScalaVersion) ||
     inMemory.exists(_.needsScalaVersion) ||
     resourceDirs.exists(_.needsScalaVersion) ||
     buildOptions.exists(_.needsScalaVersion)
+}
 
+/** Information gathered from preprocessing command inputs - sources and build options from using
+  * directives
+  *
+  * @param paths
+  *   paths and realtive paths to sources on disk, wrapped in their build requirements
+  * @param inMemory
+  *   in memory sources (e.g. snippets and wrapped scripts) wrapped in their build requirements
+  * @param defaultMainClass
+  * @param resourceDirs
+  * @param buildOptions
+  *   build options from sources
+  */
+final case class CrossSources(
+  paths: Seq[WithBuildRequirements[(os.Path, os.RelPath)]],
+  inMemory: Seq[WithBuildRequirements[Sources.InMemory]],
+  defaultMainClass: Option[String],
+  resourceDirs: Seq[WithBuildRequirements[os.Path]],
+  buildOptions: Seq[WithBuildRequirements[BuildOptions]]
+) extends UnwrappedCrossSources(
+      paths,
+      inMemory,
+      defaultMainClass,
+      resourceDirs,
+      buildOptions,
+      Nil
+    ) {
   def scopedSources(baseOptions: BuildOptions): Either[BuildException, ScopedSources] = either {
 
     val sharedOptions0 = sharedOptions(baseOptions)
@@ -114,34 +182,6 @@ final case class CrossSources(
       crossSources0.buildOptions.map(_.scopedValue(defaultScope))
     )
   }
-
-  /** For all unwrapped script sources contained in this object wrap them according to provided
-    * BuildOptions
-    *
-    * @param buildOptions
-    *   options used to choose the script wrapper
-    * @return
-    *   this with all the script code wrapped
-    */
-  def withWrappedScripts(buildOptions: BuildOptions): CrossSources =
-    copy(
-      inMemory = inMemory.map {
-        case WithBuildRequirements(requirements, source) if source.wrapScriptFunOpt.isDefined =>
-          val wrapScriptFun                = source.wrapScriptFunOpt.get
-          val codeWrapper                  = ScriptPreprocessor.getScriptWrapper(buildOptions)
-          val (wrappedCode, topWrapperLen) = wrapScriptFun(codeWrapper)
-
-          WithBuildRequirements(
-            requirements,
-            source.copy(
-              generatedContent = wrappedCode,
-              topWrapperLen = topWrapperLen,
-              wrapScriptFunOpt = None
-            )
-          )
-        case p => p
-      }
-    )
 }
 
 object CrossSources {
@@ -168,7 +208,7 @@ object CrossSources {
     suppressWarningOptions: SuppressWarningOptions,
     exclude: Seq[Positioned[String]] = Nil,
     maybeRecoverOnError: BuildException => Option[BuildException] = e => Some(e)
-  )(using ScalaCliInvokeData): Either[BuildException, (CrossSources, Inputs)] = either {
+  )(using ScalaCliInvokeData): Either[BuildException, (UnwrappedCrossSources, Inputs)] = either {
 
     def preprocessSources(elems: Seq[SingleElement])
       : Either[BuildException, Seq[PreprocessedSource]] =
@@ -286,7 +326,17 @@ object CrossSources {
           val baseReqs0 = baseReqs(m.scopePath)
           WithBuildRequirements(
             m.requirements.fold(baseReqs0)(_ orElse baseReqs0),
-            Sources.InMemory(m.originalPath, m.relPath, m.code, m.ignoreLen, m.wrapScriptFunOpt)
+            Sources.InMemory(m.originalPath, m.relPath, m.code, m.ignoreLen)
+          ) -> m.directivesPositions
+      }
+    val unwrappedScriptsWithDirectivePositions
+      : Seq[(WithBuildRequirements[Sources.UnwrappedScript], Option[DirectivesPositions])] =
+      preprocessedSources.collect {
+        case m: PreprocessedSource.UnwrappedScript =>
+          val baseReqs0 = baseReqs(m.scopePath)
+          WithBuildRequirements(
+            m.requirements.fold(baseReqs0)(_ orElse baseReqs0),
+            Sources.UnwrappedScript(m.originalPath, m.relPath, m.wrapScriptFun)
           ) -> m.directivesPositions
       }
 
@@ -298,14 +348,20 @@ object CrossSources {
     )
 
     lazy val allPathsWithDirectivesByScope: Map[Scope, Seq[(os.Path, DirectivesPositions)]] =
-      (pathsWithDirectivePositions ++ inMemoryWithDirectivePositions)
+      (pathsWithDirectivePositions ++
+        inMemoryWithDirectivePositions ++
+        unwrappedScriptsWithDirectivePositions)
         .flatMap { (withBuildRequirements, directivesPositions) =>
           val scope = withBuildRequirements.scopedValue(Scope.Main).scope
           val path: os.Path = withBuildRequirements.value match
             case im: Sources.InMemory =>
               im.originalPath match
                 case Right((_, p: os.Path)) => p
                 case _                      => inputs.workspace / im.generatedRelPath
+            case us: Sources.UnwrappedScript =>
+              us.originalPath match
+                case Right((_, p: os.Path)) => p
+                case _                      => inputs.workspace / us.generatedRelPath
             case (p: os.Path, _) => p
           directivesPositions.map((path, scope, _))
         }
@@ -333,9 +389,20 @@ object CrossSources {
         }
     }
 
-    val paths    = pathsWithDirectivePositions.map(_._1)
-    val inMemory = inMemoryWithDirectivePositions.map(_._1)
-    (CrossSources(paths, inMemory, defaultMainClassOpt, resourceDirs, buildOptions), allInputs)
+    val paths            = pathsWithDirectivePositions.map(_._1)
+    val inMemory         = inMemoryWithDirectivePositions.map(_._1)
+    val unwrappedScripts = unwrappedScriptsWithDirectivePositions.map(_._1)
+    (
+      UnwrappedCrossSources(
+        paths,
+        inMemory,
+        defaultMainClassOpt,
+        resourceDirs,
+        buildOptions,
+        unwrappedScripts
+      ),
+      allInputs
+    )
   }
 
   private def resolveInputsFromSources(sources: Seq[Positioned[os.Path]], enableMarkdown: Boolean) =

diff --git a/modules/build/src/main/scala/scala/build/Sources.scala b/modules/build/src/main/scala/scala/build/Sources.scala
@@ -70,10 +70,20 @@ object Sources {
     originalPath: Either[String, (os.SubPath, os.Path)],
     generatedRelPath: os.RelPath,
     generatedContent: String,
-    topWrapperLen: Int,
-    wrapScriptFunOpt: Option[CodeWrapper => (String, Int)] = None
+    topWrapperLen: Int
   )
 
+  final case class UnwrappedScript(
+    originalPath: Either[String, (os.SubPath, os.Path)],
+    generatedRelPath: os.RelPath,
+    wrapScriptFun: CodeWrapper => (String, Int)
+  ) {
+    def wrap(wrapper: CodeWrapper): InMemory = {
+      val (content, topWrapperLen) = wrapScriptFun(wrapper)
+      InMemory(originalPath, generatedRelPath, content, topWrapperLen)
+    }
+  }
+
   /** The default preprocessor list.
     *
     * @param codeWrapper

diff --git a/modules/build/src/main/scala/scala/build/preprocessing/PreprocessedSource.scala b/modules/build/src/main/scala/scala/build/preprocessing/PreprocessedSource.scala
@@ -39,12 +39,25 @@ object PreprocessedSource {
     scopedRequirements: Seq[Scoped[BuildRequirements]],
     mainClassOpt: Option[String],
     scopePath: ScopePath,
-    directivesPositions: Option[DirectivesPositions],
-    wrapScriptFunOpt: Option[CodeWrapper => (String, Int)] = None
+    directivesPositions: Option[DirectivesPositions]
   ) extends PreprocessedSource {
     def reportingPath: Either[String, os.Path] =
       originalPath.map(_._2)
   }
+
+  final case class UnwrappedScript(
+    originalPath: Either[String, (os.SubPath, os.Path)],
+    relPath: os.RelPath,
+    options: Option[BuildOptions],
+    optionsWithTargetRequirements: List[WithBuildRequirements[BuildOptions]],
+    requirements: Option[BuildRequirements],
+    scopedRequirements: Seq[Scoped[BuildRequirements]],
+    mainClassOpt: Option[String],
+    scopePath: ScopePath,
+    directivesPositions: Option[DirectivesPositions],
+    wrapScriptFun: CodeWrapper => (String, Int)
+  ) extends PreprocessedSource
+
   final case class NoSourceCode(
     options: Option[BuildOptions],
     optionsWithTargetRequirements: List[WithBuildRequirements[BuildOptions]],

diff --git a/modules/build/src/main/scala/scala/build/preprocessing/ScriptPreprocessor.scala b/modules/build/src/main/scala/scala/build/preprocessing/ScriptPreprocessor.scala
@@ -76,66 +76,65 @@ case object ScriptPreprocessor extends Preprocessor {
     maybeRecoverOnError: BuildException => Option[BuildException],
     allowRestrictedFeatures: Boolean,
     suppressWarningOptions: SuppressWarningOptions
-  )(using ScalaCliInvokeData): Either[BuildException, List[PreprocessedSource.InMemory]] = either {
-
-    val (contentIgnoredSheBangLines, _) = SheBang.ignoreSheBangLines(content)
-
-    val (pkg, wrapper) = AmmUtil.pathToPackageWrapper(subPath)
-
-    val processingOutput: ProcessingOutput =
-      value(ScalaPreprocessor.process(
-        contentIgnoredSheBangLines,
-        reportingPath,
-        scopePath / os.up,
-        logger,
-        maybeRecoverOnError,
-        allowRestrictedFeatures,
-        suppressWarningOptions
-      ))
-        .getOrElse(ProcessingOutput.empty)
-
-    val scriptCode = processingOutput.updatedContent.getOrElse(contentIgnoredSheBangLines)
-    // try to match in multiline mode, don't match comment lines starting with '//'
-    val containsMainAnnot = "(?m)^(?!//).*@main.*".r.findFirstIn(scriptCode).isDefined
-
-    val wrapScriptFun = (cw: CodeWrapper) => {
-      if (containsMainAnnot) logger.diagnostic(
-        cw match {
-          case _: ObjectCodeWrapper.type =>
-            WarningMessages.mainAnnotationNotSupported( /* annotationIgnored */ true)
-          case _ => WarningMessages.mainAnnotationNotSupported( /* annotationIgnored */ false)
-        }
-      )
-
-      val (code, topWrapperLen, _) = cw.wrapCode(
-        pkg,
-        wrapper,
-        scriptCode,
-        inputArgPath.getOrElse(subPath.last)
+  )(using ScalaCliInvokeData): Either[BuildException, List[PreprocessedSource.UnwrappedScript]] =
+    either {
+
+      val (contentIgnoredSheBangLines, _) = SheBang.ignoreSheBangLines(content)
+
+      val (pkg, wrapper) = AmmUtil.pathToPackageWrapper(subPath)
+
+      val processingOutput: ProcessingOutput =
+        value(ScalaPreprocessor.process(
+          contentIgnoredSheBangLines,
+          reportingPath,
+          scopePath / os.up,
+          logger,
+          maybeRecoverOnError,
+          allowRestrictedFeatures,
+          suppressWarningOptions
+        ))
+          .getOrElse(ProcessingOutput.empty)
+
+      val scriptCode = processingOutput.updatedContent.getOrElse(contentIgnoredSheBangLines)
+      // try to match in multiline mode, don't match comment lines starting with '//'
+      val containsMainAnnot = "(?m)^(?!//).*@main.*".r.findFirstIn(scriptCode).isDefined
+
+      val wrapScriptFun = (cw: CodeWrapper) => {
+        if (containsMainAnnot) logger.diagnostic(
+          cw match {
+            case _: ObjectCodeWrapper.type =>
+              WarningMessages.mainAnnotationNotSupported( /* annotationIgnored */ true)
+            case _ => WarningMessages.mainAnnotationNotSupported( /* annotationIgnored */ false)
+          }
+        )
+
+        val (code, topWrapperLen, _) = cw.wrapCode(
+          pkg,
+          wrapper,
+          scriptCode,
+          inputArgPath.getOrElse(subPath.last)
+        )
+        (code, topWrapperLen)
+      }
+
+      val className = (pkg :+ wrapper).map(_.raw).mkString(".")
+      val relPath   = os.rel / (subPath / os.up) / s"${subPath.last.stripSuffix(".sc")}.scala"
+
+      val file = PreprocessedSource.UnwrappedScript(
+        originalPath = reportingPath.map((subPath, _)),
+        relPath = relPath,
+        options = Some(processingOutput.opts),
+        optionsWithTargetRequirements = processingOutput.optsWithReqs,
+        requirements = Some(processingOutput.globalReqs),
+        scopedRequirements = processingOutput.scopedReqs,
+        mainClassOpt = Some(CodeWrapper.mainClassObject(Name(className)).backticked),
+        scopePath = scopePath,
+        directivesPositions = processingOutput.directivesPositions,
+        wrapScriptFun = wrapScriptFun
       )
-      (code, topWrapperLen)
+      List(file)
     }
 
-    val className = (pkg :+ wrapper).map(_.raw).mkString(".")
-    val relPath   = os.rel / (subPath / os.up) / s"${subPath.last.stripSuffix(".sc")}.scala"
-
-    val file = PreprocessedSource.InMemory(
-      originalPath = reportingPath.map((subPath, _)),
-      relPath = relPath,
-      code = "", // code is captured in wrapScriptFun's closure
-      ignoreLen = 0,
-      options = Some(processingOutput.opts),
-      optionsWithTargetRequirements = processingOutput.optsWithReqs,
-      requirements = Some(processingOutput.globalReqs),
-      scopedRequirements = processingOutput.scopedReqs,
-      mainClassOpt = Some(CodeWrapper.mainClassObject(Name(className)).backticked),
-      scopePath = scopePath,
-      directivesPositions = processingOutput.directivesPositions,
-      wrapScriptFunOpt = Some(wrapScriptFun)
-    )
-    List(file)
-  }
-
   /** Get correct script wrapper depending on the platform and version of Scala. For Scala 2 or
     * Platform JS use [[ObjectCodeWrapper]]. Otherwise - for Scala 3 on JVM or Native use
     * [[ClassCodeWrapper]].