Add import support to main-kts, with test

libraries/tools/kotlin-main-kts-test/test/org/jetbrains/kotlin/mainKts/test/mainKtsTest.kt

@@ -7,11 +7,10 @@ package org.jetbrains.kotlin.mainKts.test
 import org.jetbrains.kotlin.mainKts.MainKtsScript
 import org.junit.Assert
 import org.junit.Test
+import java.io.ByteArrayOutputStream
 import java.io.File
-import kotlin.script.experimental.api.EvaluationResult
-import kotlin.script.experimental.api.ResultWithDiagnostics
-import kotlin.script.experimental.api.ScriptEvaluationConfiguration
-import kotlin.script.experimental.api.constructorArgs
+import java.io.PrintStream
+import kotlin.script.experimental.api.*
 import kotlin.script.experimental.host.toScriptSource
 import kotlin.script.experimental.jvmhost.BasicJvmScriptingHost
 import kotlin.script.experimental.jvmhost.createJvmCompilationConfigurationFromTemplate
@@ -22,6 +21,7 @@ fun evalFile(scriptFile: File): ResultWithDiagnostics<EvaluationResult> {
 
     val evaluationEnv = ScriptEvaluationConfiguration {
         constructorArgs(emptyArray<String>())
+        enableScriptsInstancesSharing()
     }
 
     return BasicJvmScriptingHost().eval(scriptFile.toScriptSource(), scriptDefinition, evaluationEnv)
@@ -60,4 +60,32 @@ class MainKtsTest {
                     ":\n  ${res.reports.joinToString("\n  ") { it.message + if (it.exception == null) "" else ": ${it.exception}" }}",
             res is ResultWithDiagnostics.Failure && res.reports.any { it.message.contains("Unknown set of arguments to maven resolver: abracadabra") })
     }
-}
+
+    @Test
+    fun testImport() {
+
+        val out = captureOut {
+            val res = evalFile(File("testData/import-test.main.kts"))
+
+            Assert.assertTrue(
+                "test failed:\n  ${res.reports.joinToString("\n  ") { it.message + if (it.exception == null) "" else ": ${it.exception}" }}",
+                res is ResultWithDiagnostics.Success
+            )
+        }.lines()
+
+        Assert.assertEquals(listOf("Hi from common", "Hi from middle", "sharedVar == 5"), out)
+    }
+}
+
+private fun captureOut(body: () -> Unit): String {
+    val outStream = ByteArrayOutputStream()
+    val prevOut = System.out
+    System.setOut(PrintStream(outStream))
+    try {
+        body()
+    } finally {
+        System.out.flush()
+        System.setOut(prevOut)
+    }
+    return outStream.toString().trim()
+}

libraries/tools/kotlin-main-kts-test/testData/import-common.main.kts

@@ -0,0 +1,4 @@
+
+var sharedVar = 2
+
+println("Hi from common")

libraries/tools/kotlin-main-kts-test/testData/import-middle.main.kts

@@ -0,0 +1,6 @@
+
+@file:Import("import-common.main.kts")
+
+sharedVar *= 2
+
+println("Hi from middle")

libraries/tools/kotlin-main-kts-test/testData/import-test.main.kts

@@ -0,0 +1,7 @@
+
+@file:Import("import-common.main.kts")
+@file:Import("import-middle.main.kts")
+
+sharedVar = sharedVar + 1
+
+println("sharedVar == $sharedVar")

libraries/tools/kotlin-main-kts/src/org/jetbrains/kotlin/mainKts/scriptDef.kt

@@ -8,11 +8,13 @@ package org.jetbrains.kotlin.mainKts
 import org.jetbrains.kotlin.mainKts.impl.FilesAndIvyResolver
 import org.jetbrains.kotlin.script.util.DependsOn
 import org.jetbrains.kotlin.script.util.Repository
+import org.jetbrains.kotlin.script.util.Import
 import java.io.File
 import kotlin.script.dependencies.ScriptContents
 import kotlin.script.dependencies.ScriptDependenciesResolver
 import kotlin.script.experimental.annotations.KotlinScript
 import kotlin.script.experimental.api.*
+import kotlin.script.experimental.host.FileScriptSource
 import kotlin.script.experimental.jvm.JvmDependency
 import kotlin.script.experimental.jvm.compat.mapLegacyDiagnosticSeverity
 import kotlin.script.experimental.jvm.compat.mapLegacyScriptPosition
@@ -25,12 +27,12 @@ abstract class MainKtsScript(val args: Array<String>)
 
 object MainKtsScriptDefinition : ScriptCompilationConfiguration(
     {
-        defaultImports(DependsOn::class, Repository::class)
+        defaultImports(DependsOn::class, Repository::class, Import::class)
         jvm {
             dependenciesFromClassContext(MainKtsScriptDefinition::class, "kotlin-main-kts")
         }
         refineConfiguration {
-            onAnnotations(DependsOn::class, Repository::class, handler = MainKtsConfigurator())
+            onAnnotations(DependsOn::class, Repository::class, Import::class, handler = MainKtsConfigurator())
         }
         ide {
             acceptedLocations(ScriptAcceptedLocation.Everywhere)
@@ -41,14 +43,8 @@ class MainKtsConfigurator : RefineScriptCompilationConfigurationHandler {
     private val resolver = FilesAndIvyResolver()
 
     override operator fun invoke(context: ScriptConfigurationRefinementContext): ResultWithDiagnostics<ScriptCompilationConfiguration> {
-        val annotations = context.collectedData?.get(ScriptCollectedData.foundAnnotations)?.takeIf { it.isNotEmpty() }
-            ?: return context.compilationConfiguration.asSuccess()
-        val scriptContents = object : ScriptContents {
-            override val annotations: Iterable<Annotation> = annotations
-            override val file: File? = null
-            override val text: CharSequence? = null
-        }
         val diagnostics = arrayListOf<ScriptDiagnostic>()
+
         fun report(severity: ScriptDependenciesResolver.ReportSeverity, message: String, position: ScriptContents.Position?) {
             diagnostics.add(
                 ScriptDiagnostic(
@@ -59,18 +55,33 @@ class MainKtsConfigurator : RefineScriptCompilationConfigurationHandler {
                 )
             )
         }
-        return try {
-            val newDepsFromResolver = resolver.resolve(scriptContents, emptyMap(), ::report, null).get()
-                ?: return context.compilationConfiguration.asSuccess(diagnostics) // TODO: failure
-            val resolvedClasspath = newDepsFromResolver.classpath.toList().takeIf { it.isNotEmpty() }
-                ?: return context.compilationConfiguration.asSuccess(diagnostics) // TODO: failure
 
-            ScriptCompilationConfiguration(context.compilationConfiguration) {
-                dependencies.append(JvmDependency(resolvedClasspath))
-            }.asSuccess(diagnostics)
-        } catch (e: Throwable) {
-            ResultWithDiagnostics.Failure(*diagnostics.toTypedArray(), e.asDiagnostics(path = context.script.locationId))
+        val annotations = context.collectedData?.get(ScriptCollectedData.foundAnnotations)?.takeIf { it.isNotEmpty() }
+            ?: return context.compilationConfiguration.asSuccess()
+
+        val scriptBaseDir = (context.script as? FileScriptSource)?.file?.parentFile
+        val importedSources = annotations.flatMap {
+            (it as? Import)?.paths?.map { sourceName ->
+                FileScriptSource(scriptBaseDir?.resolve(sourceName) ?: File(sourceName))
+            } ?: emptyList()
         }
+
+        val resolvedClassPath = try {
+            val scriptContents = object : ScriptContents {
+                override val annotations: Iterable<Annotation> = annotations.filter { it is DependsOn || it is Repository }
+                override val file: File? = null
+                override val text: CharSequence? = null
+            }
+            resolver.resolve(scriptContents, emptyMap(), ::report, null).get()?.classpath?.toList()
+            // TODO: add diagnostics
+        } catch (e: Throwable) {
+            return ResultWithDiagnostics.Failure(*diagnostics.toTypedArray(), e.asDiagnostics(path = context.script.locationId))
+        }
+
+        return ScriptCompilationConfiguration(context.compilationConfiguration) {
+            if (resolvedClassPath?.isNotEmpty() == true) dependencies.append(JvmDependency(resolvedClassPath))
+            if (importedSources.isNotEmpty()) importScripts.append(importedSources)
+        }.asSuccess(diagnostics)
     }
 }
 

libraries/tools/kotlin-script-util/src/main/kotlin/org/jetbrains/kotlin/script/util/annotations.kt

@@ -30,3 +30,10 @@ annotation class DependsOn(val value: String = "", val groupId: String = "", val
 @Retention(AnnotationRetention.SOURCE)
 annotation class Repository(val value: String = "", val id: String = "", val url: String = "")
 
+/**
+ * Import other script(s)
+ */
+@Target(AnnotationTarget.FILE)
+@Repeatable
+@Retention(AnnotationRetention.SOURCE)
+annotation class Import(vararg val paths: String)