Register custom file extension for scripts loaded from script template

idea/idea-core/src/org/jetbrains/kotlin/idea/core/script/ScriptDefinitionsManager.kt

@@ -18,10 +18,13 @@ package org.jetbrains.kotlin.idea.core.script
 
 import com.intellij.ide.projectView.impl.ProjectRootsUtil.isInTestSource
 import com.intellij.ide.scratch.ScratchFileService
+import com.intellij.openapi.application.ApplicationManager
+import com.intellij.openapi.application.runWriteAction
 import com.intellij.openapi.components.ServiceManager
 import com.intellij.openapi.diagnostic.Logger
 import com.intellij.openapi.extensions.ExtensionPointName
 import com.intellij.openapi.extensions.Extensions
+import com.intellij.openapi.fileTypes.FileTypeManager
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.projectRoots.JavaSdk
 import com.intellij.openapi.projectRoots.ProjectJdkTable
@@ -29,21 +32,18 @@ import com.intellij.openapi.projectRoots.ex.PathUtilEx
 import com.intellij.openapi.roots.ProjectRootManager
 import com.intellij.openapi.vfs.VfsUtil
 import com.intellij.openapi.vfs.VirtualFile
+import org.jetbrains.kotlin.idea.KotlinFileType
 import org.jetbrains.kotlin.idea.caches.project.SdkInfo
 import org.jetbrains.kotlin.idea.caches.project.getScriptRelatedModuleInfo
-import org.jetbrains.kotlin.script.*
-import org.jetbrains.kotlin.scripting.compiler.plugin.KotlinScriptDefinitionAdapterFromNewAPI
 import org.jetbrains.kotlin.idea.util.ProjectRootsUtil.isInContent
 import org.jetbrains.kotlin.psi.KtFile
-import org.jetbrains.kotlin.script.KotlinScriptDefinition
+import org.jetbrains.kotlin.script.*
-import org.jetbrains.kotlin.script.KotlinScriptDefinitionFromAnnotatedTemplate
+import org.jetbrains.kotlin.scripting.compiler.plugin.KotlinScriptDefinitionAdapterFromNewAPI
-import org.jetbrains.kotlin.script.ScriptDefinitionProvider
+import org.jetbrains.kotlin.scripting.compiler.plugin.KotlinScriptDefinitionAdapterFromNewAPIBase
-import org.jetbrains.kotlin.script.ScriptTemplatesProvider
 import org.jetbrains.kotlin.utils.PathUtil
 import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstanceOrNull
 import org.jetbrains.kotlin.utils.addToStdlib.flattenTo
 import java.io.File
-import java.lang.reflect.InvocationTargetException
 import java.net.URLClassLoader
 import kotlin.concurrent.write
 import kotlin.script.dependencies.Environment
@@ -133,6 +133,25 @@ class ScriptDefinitionsManager(private val project: Project) : LazyScriptDefinit
     private fun updateDefinitions() {
         assert(lock.isWriteLocked) { "updateDefinitions should only be called under the write lock" }
         definitions = definitionsByContributor.values.flattenTo(mutableListOf()).asSequence()
+
+        // Register new file extensions
+        val fileTypeManager = FileTypeManager.getInstance()
+        val extensions = definitions?.mapNotNull { definition ->
+            (definition as? KotlinScriptDefinitionAdapterFromNewAPIBase)
+                ?.scriptFileExtensionWithDot?.removePrefix(".")
+                ?.takeIf { fileTypeManager.getFileTypeByExtension(it) != KotlinFileType.INSTANCE }
+        }?.toList()
+
+        if (extensions?.isNotEmpty() == true) {
+            ApplicationManager.getApplication().invokeLater {
+                runWriteAction {
+                    extensions.forEach {
+                        fileTypeManager.associateExtension(KotlinFileType.INSTANCE, it)
+                    }
+                }
+            }
+        }
+
         clearCache()
         // TODO: clear by script type/definition
         ServiceManager.getService(project, ScriptDependenciesCache::class.java).clear()

idea/testData/script/definition/highlighting/customExtension/mainModule/script.mykts

@@ -0,0 +1 @@
+1 <error>+</error> "a"

idea/testData/script/definition/highlighting/customExtension/mainModule/template/template.kt

@@ -0,0 +1,14 @@
+package custom.scriptDefinition
+
+import java.io.File
+import kotlin.script.dependencies.*
+import kotlin.script.experimental.dependencies.*
+import kotlin.script.experimental.location.*
+import kotlin.script.templates.ScriptTemplateDefinition
+import kotlin.script.experimental.annotations.KotlinScript
+import kotlin.script.experimental.annotations.KotlinScriptEvaluator
+import kotlin.script.experimental.annotations.KotlinScriptFileExtension
+
+@KotlinScript("Kotlin Script with custom extension")
+@KotlinScriptFileExtension("mykts")
+open class Template

idea/tests/org/jetbrains/kotlin/idea/caches/resolve/IdeaModuleInfoTest.kt

@@ -27,6 +27,7 @@ import com.intellij.openapi.roots.libraries.Library
 import com.intellij.openapi.vfs.LocalFileSystem
 import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.openapi.vfs.newvfs.impl.VfsRootAccess
+import com.intellij.psi.PsiManager
 import com.intellij.testFramework.ModuleTestCase
 import com.intellij.testFramework.PlatformTestCase
 import com.intellij.testFramework.PsiTestUtil
@@ -429,7 +430,9 @@ class IdeaModuleInfoTest : ModuleTestCase() {
     ) = ModuleRootModificationUtil.addDependency(this, other, dependencyScope, exported)
 
     private val VirtualFile.moduleInfo: IdeaModuleInfo
-        get() = getModuleInfoByVirtualFile(project, this)!!
+        get() {
+            return PsiManager.getInstance(project).findFile(this)!!.getModuleInfo()
+        }
 
     private val Module.production: ModuleProductionSourceInfo
         get() = productionSourceInfo()!!

idea/tests/org/jetbrains/kotlin/idea/script/AbstractScriptConfigurationTest.kt

@@ -31,6 +31,7 @@ import com.intellij.testFramework.PlatformTestCase
 import com.intellij.testFramework.PlatformTestUtil
 import com.intellij.testFramework.PsiTestUtil
 import com.intellij.testFramework.exceptionCases.AbstractExceptionCase
+import com.intellij.util.ui.UIUtil
 import org.jetbrains.kotlin.codegen.forTestCompile.ForTestCompileRuntime
 import org.jetbrains.kotlin.idea.completion.test.KotlinCompletionTestCase
 import org.jetbrains.kotlin.idea.core.script.ScriptDefinitionContributor
@@ -46,6 +47,9 @@ import org.jetbrains.kotlin.test.util.addDependency
 import org.jetbrains.kotlin.test.util.projectLibrary
 import org.jetbrains.kotlin.test.util.renderAsGotoImplementation
 import org.jetbrains.kotlin.utils.PathUtil
+import org.jetbrains.kotlin.utils.PathUtil.KOTLIN_JAVA_SCRIPT_RUNTIME_JAR
+import org.jetbrains.kotlin.utils.PathUtil.KOTLIN_SCRIPTING_COMMON_JAR
+import org.jetbrains.kotlin.utils.PathUtil.KOTLIN_SCRIPTING_JVM_JAR
 import org.junit.Assert
 import org.junit.ComparisonFailure
 import java.io.File
@@ -63,7 +67,8 @@ abstract class AbstractScriptConfigurationHighlightingTest : AbstractScriptConfi
         checkHighlighting(
             editor,
             InTextDirectivesUtils.isDirectiveDefined(file.text, "// CHECK_WARNINGS"),
-            InTextDirectivesUtils.isDirectiveDefined(file.text, "// CHECK_INFOS"))
+            InTextDirectivesUtils.isDirectiveDefined(file.text, "// CHECK_INFOS")
+        )
     }
 
     fun doComplexTest(path: String) {
@@ -131,8 +136,13 @@ abstract class AbstractScriptConfigurationTest : KotlinCompletionTestCase() {
         // do not create default module
     }
 
-    private fun findMainScript(testDir: String) = File(testDir).walkTopDown().find { it.name == SCRIPT_NAME }
+    private fun findMainScript(testDir: String): File {
+        val scriptFile = File(testDir).walkTopDown().find { it.name == SCRIPT_NAME }
+        if (scriptFile != null) return scriptFile
+
+        return File(testDir).walkTopDown().singleOrNull { it.name.contains("script") }
             ?: error("Couldn't find $SCRIPT_NAME file in $testDir")
+    }
 
     protected fun configureScriptFile(path: String) {
         val mainScriptFile = findMainScript(path)
@@ -217,9 +227,9 @@ abstract class AbstractScriptConfigurationTest : KotlinCompletionTestCase() {
 
     private fun defaultEnvironment(path: String): Map<String, File?> {
         val templateOutDir = File("${path}template").takeIf { it.isDirectory }?.let {
-            compileLibToDir(it, PathUtil.kotlinPathsForDistDirectory.scriptRuntimePath.path)
+            compileLibToDir(it, *scriptClasspath())
         } ?: File("idea/testData/script/definition/defaultTemplate").takeIf { it.isDirectory }?.let {
-            compileLibToDir(it, PathUtil.kotlinPathsForDistDirectory.scriptRuntimePath.path)
+            compileLibToDir(it, *scriptClasspath())
         }
 
         val libSrcDir = File("${path}lib").takeIf { it.isDirectory }
@@ -235,14 +245,24 @@ abstract class AbstractScriptConfigurationTest : KotlinCompletionTestCase() {
         )
     }
 
+    private fun scriptClasspath(): Array<String> {
+        return with(PathUtil.kotlinPathsForDistDirectory) {
+            arrayOf(
+                File(libPath, KOTLIN_JAVA_SCRIPT_RUNTIME_JAR).path,
+                File(libPath, KOTLIN_SCRIPTING_COMMON_JAR).path,
+                File(libPath, KOTLIN_SCRIPTING_JVM_JAR).path
+            )
+        }
+    }
+
     private fun createFileAndSyncDependencies(scriptFile: File) {
         var script: VirtualFile? = null
         if (module != null) {
-            script = module.moduleFile?.parent?.findChild(SCRIPT_NAME)
+            script = module.moduleFile?.parent?.findChild(scriptFile.name)
         }
 
         if (script == null) {
-            val target = File(project.basePath, SCRIPT_NAME)
+            val target = File(project.basePath, scriptFile.name)
             scriptFile.copyTo(target)
             script = LocalFileSystem.getInstance().findFileByPath(target.path)
         }
@@ -287,6 +307,9 @@ abstract class AbstractScriptConfigurationTest : KotlinCompletionTestCase() {
             provider,
             testRootDisposable
         )
+
         ScriptDefinitionsManager.getInstance(project).reloadScriptDefinitions()
+
+        UIUtil.dispatchAllInvocationEvents()
     }
 }

idea/tests/org/jetbrains/kotlin/idea/script/ScriptConfigurationHighlightingTestGenerated.java

@@ -56,6 +56,11 @@ public class ScriptConfigurationHighlightingTestGenerated extends AbstractScript
             runTest("idea/testData/script/definition/highlighting/customBaseClass/");
         }
 
+        @TestMetadata("customExtension")
+        public void testCustomExtension() throws Exception {
+            runTest("idea/testData/script/definition/highlighting/customExtension/");
+        }
+
         @TestMetadata("customLibrary")
         public void testCustomLibrary() throws Exception {
             runTest("idea/testData/script/definition/highlighting/customLibrary/");

plugins/scripting/scripting-cli/src/org/jetbrains/kotlin/scripting/compiler/plugin/KotlinScriptDefinitionAdapterFromNewAPI.kt

@@ -24,7 +24,7 @@ abstract class KotlinScriptDefinitionAdapterFromNewAPIBase : KotlinScriptDefinit
 
     protected abstract val scriptDefinition: ScriptDefinition
 
-    protected abstract val scriptFileExtensionWithDot: String
+    abstract val scriptFileExtensionWithDot: String
 
     open val baseClass: KClass<*> by lazy(LazyThreadSafetyMode.PUBLICATION) {
         getScriptingClass(scriptDefinition.compilationConfigurator.defaultConfiguration[ScriptingEnvironmentProperties.baseClass])