Add possibility to get location of the script.main.kts file

#KT-48414 fixed
2021-09-19 21:53:13 +03:00
parent 7ddf83f32d
commit ca2f37f6eb
13 changed files with 177 additions and 7 deletions
@@ -7,6 +7,7 @@

 package kotlin.script.experimental.api

+import java.io.File
 import java.io.Serializable
 import kotlin.reflect.KClass
 import kotlin.script.experimental.host.ScriptingHostConfiguration
@@ -106,6 +107,17 @@ val ScriptCompilationConfigurationKeys.implicitReceivers by PropertiesCollection
 */
 val ScriptCompilationConfigurationKeys.providedProperties by PropertiesCollection.key<Map<String, KotlinType>>() // external variables

+/**
+ * Variable name that holds a {@link File} instance pointing to the location of the script file
+ */
+val ScriptCompilationConfigurationKeys.scriptFileLocationVariable by PropertiesCollection.key<String>()
+
+/**
+ * File pointing to the location of the script file. Note that in some cases it might not be possible
+ * to determine script file location properly - in this case the file is an empty file
+ */
+val ScriptCompilationConfigurationKeys.scriptFileLocation by PropertiesCollection.key<File>()
+
 /**
 * The list of import expressions that will be implicitly applied to the script body, the syntax is the same as for the "import" statement
 */
@@ -41,12 +41,14 @@ fun configureProvidedPropertiesFromJsr223Context(context: ScriptEvaluationConfig
        val engineBindings = jsr223context.getBindings(ScriptContext.ENGINE_SCOPE)
        val globalBindings = jsr223context.getBindings(ScriptContext.GLOBAL_SCOPE)
        for (prop in knownProperties) {
-            val v = when {
-                engineBindings?.containsKey(prop.key) == true -> engineBindings[prop.key]
-                globalBindings?.containsKey(prop.key) == true -> globalBindings[prop.key]
-                else -> return ResultWithDiagnostics.Failure("Property ${prop.key} is not found in the bindings".asErrorDiagnostics())
+            if (prop.key !in updatedProperties) {
+                val v = when {
+                    engineBindings?.containsKey(prop.key) == true -> engineBindings[prop.key]
+                    globalBindings?.containsKey(prop.key) == true -> globalBindings[prop.key]
+                    else -> return ResultWithDiagnostics.Failure("Property ${prop.key} is not found in the bindings".asErrorDiagnostics())
+                }
+                updatedProperties[prop.key] = v
            }
-            updatedProperties[prop.key] = v
        }
        ScriptEvaluationConfiguration(context.evaluationConfiguration) {
            providedProperties(updatedProperties)
@@ -121,6 +121,26 @@ class MainKtsIT {
        }
    }

+    @OptIn(ExperimentalPathApi::class)
+    @Test
+    fun testCacheWithFileLocation() {
+        val scriptPath = File("$TEST_DATA_ROOT/script-file-location-default.main.kts").absolutePath
+        val cache = createTempDirectory("main.kts.test")
+        val expectedTestOutput = listOf(Regex.escape(scriptPath))
+
+        try {
+            Assert.assertTrue(cache.exists() && cache.listDirectoryEntries("*.jar").isEmpty())
+            runWithKotlinRunner(scriptPath, expectedTestOutput, cacheDir = cache)
+            val cacheFile = cache.listDirectoryEntries("*.jar").firstOrNull()
+            Assert.assertTrue(cacheFile != null && cacheFile.exists())
+
+            // this run should use the cached script
+            runWithKotlinRunner(scriptPath, expectedTestOutput, cacheDir = cache)
+        } finally {
+            cache.toFile().deleteRecursively()
+        }
+    }
+
    @Test
    fun testHelloSerialization() {
        val paths = PathUtil.kotlinPathsForDistDirectory
@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.mainKts.test
 import org.jetbrains.kotlin.mainKts.COMPILED_SCRIPTS_CACHE_DIR_PROPERTY
 import org.jetbrains.kotlin.mainKts.impl.Directories
 import org.jetbrains.kotlin.mainKts.MainKtsScript
+import org.jetbrains.kotlin.mainKts.SCRIPT_FILE_LOCATION_DEFAULT_VARIABLE_NAME
 import org.jetbrains.kotlin.scripting.compiler.plugin.assertTrue
 import org.junit.Assert
 import org.junit.Assert.assertEquals
@@ -145,6 +146,65 @@ class MainKtsTest {
        Assert.assertEquals(listOf("Hi from sub", "Hi from super", "Hi from random"), out)
    }

+    @Test
+    fun testScriptFileLocationDefaultVariable() {
+        val resOk = evalFile(File("$TEST_DATA_ROOT/script-file-location-default.main.kts"))
+        assertSucceeded(resOk)
+        val resultValue = resOk.valueOrThrow().returnValue
+        assertTrue(resultValue is ResultValue.Value) { "Result value should be of type Value" }
+        val value = (resultValue as ResultValue.Value).value!!
+        assertEquals("String", value::class.simpleName)
+        val expectedPathSuffix = "libraries/tools/kotlin-main-kts-test/testData/script-file-location-default.main.kts"
+        val actualPath = (value as String).replace("\\", "/")
+        assertTrue(actualPath.endsWith(expectedPathSuffix)) { "Script file path does not end with expected path" }
+    }
+
+    @Test
+    fun testScriptFileLocationCustomizedVariable() {
+        val resOk = evalFile(File("$TEST_DATA_ROOT/script-file-location-customized.main.kts"))
+        assertSucceeded(resOk)
+        val resultValue = resOk.valueOrThrow().returnValue
+        assertTrue(resultValue is ResultValue.Value) { "Result value should be of type Value" }
+        val value = (resultValue as ResultValue.Value).value!!
+        assertEquals("String", value::class.simpleName)
+        val expectedPathSuffix = "libraries/tools/kotlin-main-kts-test/testData/script-file-location-customized.main.kts"
+        val actualPath = (value as String).replace("\\", "/")
+        assertTrue(actualPath.endsWith(expectedPathSuffix)) { "Script file path does not end with expected path" }
+    }
+
+    @Test
+    fun testScriptFileLocationWithImportedScript() {
+        val resOk = evalFile(File("$TEST_DATA_ROOT/script-file-location-with-imported-file.main.kts"))
+        assertSucceeded(resOk)
+        val resultValue = resOk.valueOrThrow().returnValue
+        assertTrue(resultValue is ResultValue.Value) { "Result value should be of type Value" }
+        val value = (resultValue as ResultValue.Value).value!!
+        assertEquals("Array", value::class.simpleName)
+        val expectedSelfPathSuffix = "libraries/tools/kotlin-main-kts-test/testData/script-file-location-with-imported-file.main.kts"
+        val expectedImportedPathSuffix = "libraries/tools/kotlin-main-kts-test/testData/script-file-location-helper-imported-file.main.kts"
+        val actualPathSelf = (value as Array<*>)[0].toString().replace("\\", "/")
+        val actualPathImported = value[1].toString().replace("\\", "/")
+        assertTrue(actualPathSelf.endsWith(expectedSelfPathSuffix)) { "Script file path does not end with expected path" }
+        assertTrue(actualPathImported.endsWith(expectedImportedPathSuffix)) { "Script file path does not end with expected path" }
+    }
+
+    @Test
+    fun testScriptFileLocationDefaultVariableNotAvailableIfScriptFileVariableCustomized() {
+        val resFailed = evalFile(File("$TEST_DATA_ROOT/script-file-location-customized-default-not-available.main.kts"))
+        assertFailed("Unresolved reference: $SCRIPT_FILE_LOCATION_DEFAULT_VARIABLE_NAME", resFailed)
+    }
+
+    @Test
+    fun testScriptFileLocationDefaultVariableRedefinition() {
+        val resOk = evalFile(File("$TEST_DATA_ROOT/script-file-location-redefine-variable.kts"))
+        assertSucceeded(resOk)
+        val resultValue = resOk.valueOrThrow().returnValue
+        assertTrue(resultValue is ResultValue.Value) { "Result value should be of type Value" }
+        val value = (resultValue as ResultValue.Value).value!!
+        assertEquals("String", value::class.simpleName)
+        assertEquals("success", value)
+    }
+
    private fun assertIsJava6Bytecode(res: ResultWithDiagnostics<EvaluationResult>) {
        val scriptClassResource = res.valueOrThrow().returnValue.scriptClass!!.java.run {
            getResource("$simpleName.class")
@@ -0,0 +1,4 @@
+
+@file:ScriptFileLocation("scriptFileLocation")
+
+__FILE__.absolutePath
@@ -0,0 +1,4 @@
+
+@file:ScriptFileLocation("scriptFileLocation")
+
+scriptFileLocation.absolutePath
@@ -0,0 +1,2 @@
+
+__FILE__.absolutePath
@@ -0,0 +1,5 @@
+import java.io.File
+
+fun getDependentScriptFile(): File {
+    return __FILE__
+}
@@ -0,0 +1,3 @@
+
+val __FILE__ = "success"
+__FILE__
@@ -0,0 +1,4 @@
+@file:Import("script-file-location-helper-imported-file.main.kts")
+
+arrayOf(__FILE__.absolutePath, getDependentScriptFile().absolutePath)
+
@@ -33,3 +33,11 @@ annotation class Import(vararg val paths: String)
@Repeatable
@Retention(AnnotationRetention.SOURCE)
 annotation class CompilerOptions(vararg val options: String)
+
+/**
+ * Option that configures the name of the variable that will hold a file pointing to the script location.
+ * If not specified, {@link [SCRIPT_FILE_LOCATION_DEFAULT_VARIABLE_NAME]} will be used as the variable name
+ */
+@Target(AnnotationTarget.FILE)
+@Retention(AnnotationRetention.SOURCE)
+annotation class ScriptFileLocation(val variable: String)
@@ -40,15 +40,18 @@ abstract class MainKtsScript(val args: Array<String>)
 const val COMPILED_SCRIPTS_CACHE_DIR_ENV_VAR = "KOTLIN_MAIN_KTS_COMPILED_SCRIPTS_CACHE_DIR"
 const val COMPILED_SCRIPTS_CACHE_DIR_PROPERTY = "kotlin.main.kts.compiled.scripts.cache.dir"
 const val COMPILED_SCRIPTS_CACHE_VERSION = 1
+const val SCRIPT_FILE_LOCATION_DEFAULT_VARIABLE_NAME = "__FILE__"

 class MainKtsScriptDefinition : ScriptCompilationConfiguration(
    {
-        defaultImports(DependsOn::class, Repository::class, Import::class, CompilerOptions::class)
+        defaultImports(DependsOn::class, Repository::class, Import::class, CompilerOptions::class, ScriptFileLocation::class)
        jvm {
            dependenciesFromClassContext(MainKtsScriptDefinition::class, "kotlin-main-kts", "kotlin-stdlib", "kotlin-reflect")
        }
        refineConfiguration {
            onAnnotations(DependsOn::class, Repository::class, Import::class, CompilerOptions::class, handler = MainKtsConfigurator())
+            onAnnotations(ScriptFileLocation::class, handler = ScriptFileLocationCustomConfigurator())
+            beforeCompiling(::configureScriptFileLocationPathVariablesForCompilation)
            beforeCompiling(::configureProvidedPropertiesFromJsr223Context)
        }
        ide {
@@ -63,6 +66,7 @@ class MainKtsScriptDefinition : ScriptCompilationConfiguration(
 object MainKtsEvaluationConfiguration : ScriptEvaluationConfiguration(
    {
        scriptsInstancesSharing(true)
+        refineConfigurationBeforeEvaluate(::configureScriptFileLocationPathVariablesForEvaluation)
        refineConfigurationBeforeEvaluate(::configureProvidedPropertiesFromJsr223Context)
        refineConfigurationBeforeEvaluate(::configureConstructorArgsFromMainArgs)
    }
@@ -90,6 +94,48 @@ class MainKtsHostConfiguration : ScriptingHostConfiguration(
    }
 )

+fun configureScriptFileLocationPathVariablesForEvaluation(context: ScriptEvaluationConfigurationRefinementContext): ResultWithDiagnostics<ScriptEvaluationConfiguration> {
+    val compilationConfiguration = context.evaluationConfiguration[ScriptEvaluationConfiguration.compilationConfiguration]
+        ?: throw RuntimeException()
+    val scriptFileLocation = compilationConfiguration[ScriptCompilationConfiguration.scriptFileLocation]
+        ?: return context.evaluationConfiguration.asSuccess()
+    val scriptFileLocationVariable = compilationConfiguration[ScriptCompilationConfiguration.scriptFileLocationVariable]
+        ?: return context.evaluationConfiguration.asSuccess()
+
+    val res = context.evaluationConfiguration.with {
+        providedProperties.put(mapOf(scriptFileLocationVariable to scriptFileLocation))
+    }
+    return res.asSuccess()
+}
+
+fun configureScriptFileLocationPathVariablesForCompilation(context: ScriptConfigurationRefinementContext): ResultWithDiagnostics<ScriptCompilationConfiguration> {
+    val scriptFile = (context.script as? FileBasedScriptSource)?.file ?: return context.compilationConfiguration.asSuccess()
+    val scriptFileLocationVariableName = context.compilationConfiguration[ScriptCompilationConfiguration.scriptFileLocationVariable]
+        ?: SCRIPT_FILE_LOCATION_DEFAULT_VARIABLE_NAME
+
+    return ScriptCompilationConfiguration(context.compilationConfiguration) {
+        providedProperties.put(mapOf(scriptFileLocationVariableName to KotlinType(File::class)))
+        scriptFileLocation.put(scriptFile)
+        scriptFileLocationVariable.put(scriptFileLocationVariableName)
+    }.asSuccess()
+}
+
+class ScriptFileLocationCustomConfigurator : RefineScriptCompilationConfigurationHandler {
+
+    override operator fun invoke(context: ScriptConfigurationRefinementContext): ResultWithDiagnostics<ScriptCompilationConfiguration> {
+
+        val scriptLocationVariable = context.collectedData?.get(ScriptCollectedData.collectedAnnotations)
+            ?.filterByAnnotationType<ScriptFileLocation>()?.firstOrNull()?.annotation?.variable
+            ?: return context.compilationConfiguration.asSuccess()
+
+        val compilationConfiguration = ScriptCompilationConfiguration(context.compilationConfiguration) {
+            scriptFileLocationVariable.put(scriptLocationVariable)
+        }
+
+        return compilationConfiguration.asSuccess()
+    }
+}
+
 fun configureConstructorArgsFromMainArgs(context: ScriptEvaluationConfigurationRefinementContext): ResultWithDiagnostics<ScriptEvaluationConfiguration> {
    val mainArgs = context.evaluationConfiguration[ScriptEvaluationConfiguration.jvm.mainArguments]
    val res = if (context.evaluationConfiguration[ScriptEvaluationConfiguration.constructorArgs] == null && mainArgs != null) {
@@ -141,7 +141,7 @@ internal fun makeCompiledScript(
                    sourceFile.declarations.firstIsInstanceOrNull<KtScript>()?.let { ktScript ->
                        makeOtherScripts(ktScript).onSuccess { otherScripts ->
                            KJvmCompiledScript(
-                                containingKtFile.virtualFilePath,
+                                sourceFile.virtualFilePath,
                                getScriptConfiguration(sourceFile),
                                ktScript.fqName.asString(),
                                null,