KT-31295: Disable module selection toolbar for .ws.kts files

^KT-31295 Fixed

idea/idea-jvm/src/org/jetbrains/kotlin/idea/scratch/ScratchFileModuleInfoProvider.kt

@@ -27,6 +27,7 @@ import org.jetbrains.kotlin.idea.core.script.scriptRelatedModuleName
 import org.jetbrains.kotlin.idea.scratch.ui.ScratchPanelListener
 import org.jetbrains.kotlin.idea.scratch.ui.ScratchTopPanel
 import org.jetbrains.kotlin.idea.util.application.runWriteAction
+import org.jetbrains.kotlin.idea.util.projectStructure.getModule
 import org.jetbrains.kotlin.parsing.KotlinParserDefinition.Companion.STD_SCRIPT_EXT
 import org.jetbrains.kotlin.parsing.KotlinParserDefinition.Companion.STD_SCRIPT_SUFFIX
 import org.jetbrains.kotlin.psi.KtFile
@@ -67,8 +68,13 @@ class ScratchFileModuleInfoProvider(val project: Project) : ProjectComponent {
                     DaemonCodeAnalyzer.getInstance(project).restart(psiFile)
                 }
 
-                val module = ktFile.virtualFile.scriptRelatedModuleName?.let { ModuleManager.getInstance(project).findModuleByName(it) }
-                if (module != null) {
+                if (file.isKotlinWorksheet) {
+                    panel.hideModuleSelector()
+
+                    val module = file.getModule(project) ?: return
+                    panel.setModule(module)
+                } else {
+                    val module = file.scriptRelatedModuleName?.let { ModuleManager.getInstance(project).findModuleByName(it) } ?: return
                     panel.setModule(module)
                 }
             }

idea/idea-jvm/src/org/jetbrains/kotlin/idea/scratch/ui/ScratchTopPanel.kt

@@ -92,6 +92,7 @@ class ScratchTopPanel private constructor(val scratchFile: ScratchFile) : JPanel
     }
 
     private val moduleChooser: ModulesComboBox
+    private val moduleChooserLabel: JLabel
     private val isReplCheckbox: JCheckBox
     private val isMakeBeforeRunCheckbox: JCheckBox
     private val isInteractiveCheckbox: JCheckBox
@@ -104,7 +105,8 @@ class ScratchTopPanel private constructor(val scratchFile: ScratchFile) : JPanel
         add(actionsToolbar.component)
 
         moduleChooser = createModuleChooser(scratchFile.project)
-        add(JLabel("Use classpath of module"))
+        moduleChooserLabel = JLabel("Use classpath of module")
+        add(moduleChooserLabel)
         add(moduleChooser)
 
         isMakeBeforeRunCheckbox = JCheckBox("Make module before Run")
@@ -160,6 +162,11 @@ class ScratchTopPanel private constructor(val scratchFile: ScratchFile) : JPanel
         moduleChooser.selectedModule = module
     }
 
+    fun hideModuleSelector() {
+        moduleChooser.isVisible = false
+        moduleChooserLabel.isVisible = false
+    }
+
     fun addModuleListener(f: (PsiFile, Module?) -> Unit) {
         moduleChooser.addActionListener {
             val selectedModule = moduleChooser.selectedModule
@@ -188,6 +195,9 @@ class ScratchTopPanel private constructor(val scratchFile: ScratchFile) : JPanel
         isInteractiveCheckbox.isSelected = isSelected
     }
 
+    @TestOnly
+    fun isModuleSelectorVisible(): Boolean = moduleChooser.isVisible && moduleChooserLabel.isVisible
+
     private fun changeMakeModuleCheckboxVisibility(isVisible: Boolean) {
         isMakeBeforeRunCheckbox.isVisible = isVisible
         moduleSeparator.isVisible = isVisible

idea/tests/org/jetbrains/kotlin/idea/scratch/AbstractScratchRunActionTest.kt

@@ -192,7 +192,9 @@ abstract class AbstractScratchRunActionTest : FileEditorManagerTestCase() {
         val (_, scratchPanel) = getEditorWithScratchPanel(myManager, myFixture.file.virtualFile)
             ?: error("Couldn't find scratch panel")
 
-        configureOptions(scratchPanel, text, myFixture.module)
+        // We want to check that correct module is selected automatically,
+        // that's why we set `module` to null so it wouldn't be changed
+        configureOptions(scratchPanel, text, null)
 
         return scratchPanel
     }

idea/tests/org/jetbrains/kotlin/idea/scratch/ScratchOptionsTest.kt

@@ -14,7 +14,7 @@ import kotlin.reflect.full.createType
 import kotlin.reflect.full.declaredMemberProperties
 
 @RunWith(JUnit3WithIdeaConfigurationRunner::class)
-class ScratchOptionsSaveTest : AbstractScratchRunActionTest() {
+class ScratchOptionsTest : AbstractScratchRunActionTest() {
 
     fun testOptionsSaveOnClosingFile() {
         val scratchPanelBeforeClosingFile = configureScratchByText("scratch_1.kts", testScratchText())
@@ -50,4 +50,27 @@ class ScratchOptionsSaveTest : AbstractScratchRunActionTest() {
             scratchPanelAfterClosingFile.scratchFile.options.isInteractiveMode
         )
     }
+
+    fun testModuleSelectionPanelIsVisibleForScratchFile() {
+        val scratchTopPanel = configureScratchByText("scratch_1.kts", testScratchText())
+
+        Assert.assertTrue("Module selector should be visible for scratches", scratchTopPanel.isModuleSelectorVisible())
+    }
+
+    fun testModuleSelectionPanelIsHiddenForWorksheetFile() {
+        val scratchTopPanel = configureWorksheetByText("worksheet.ws.kts", testScratchText())
+
+        Assert.assertFalse("Module selector should be hidden for worksheets", scratchTopPanel.isModuleSelectorVisible())
+    }
+
+    fun testCurrentModuleIsAutomaticallySelectedForWorksheetFile() {
+        val scratchTopPanel = configureWorksheetByText("worksheet.ws.kts", testScratchText())
+
+        Assert.assertEquals(
+            "Selected module should be equal to current project module for worksheets",
+            myFixture.module,
+            scratchTopPanel.getModule()
+        )
+    }
+
 }