Features/#355 module use support (#403)

feat: Basic  'module' and 'use' features support (#355)

Resolves: #355

CHANGELOG.md

@@ -18,6 +18,7 @@ Released on ...
 - TODO  (see [#NNN](https://github.com/JetBrains-Research/snakecharm/issues/NNN))
 
 ### Added
+- Support for 'module' and 'use' keywords (see [#355](https://github.com/JetBrains-Research/snakecharm/issues/355))
 - Inspection for improperly called functions (see [#148](https://github.com/JetBrains-Research/snakecharm/issues/148))
 - Ability for memorising new section name (see [#372](https://github.com/JetBrains-Research/snakecharm/issues/372))
 - Support for 'handover' section (see [#362](https://github.com/JetBrains-Research/snakecharm/issues/362))

src/main/java/com/jetbrains/snakecharm/lang/psi/elementTypes/SmkStubElementTypes.java

@@ -1,18 +1,14 @@
 package com.jetbrains.snakecharm.lang.psi.elementTypes;
 
 import com.intellij.psi.stubs.IStubElementType;
-import com.jetbrains.snakecharm.lang.psi.SmkCheckPoint;
-import com.jetbrains.snakecharm.lang.psi.SmkRule;
-import com.jetbrains.snakecharm.lang.psi.SmkSubworkflow;
-import com.jetbrains.snakecharm.lang.psi.impl.stubs.SmkCheckpointElementType;
-import com.jetbrains.snakecharm.lang.psi.impl.stubs.SmkRuleElementType;
-import com.jetbrains.snakecharm.lang.psi.impl.stubs.SmkSubworkflowElementType;
-import com.jetbrains.snakecharm.lang.psi.stubs.SmkCheckpointStub;
-import com.jetbrains.snakecharm.lang.psi.stubs.SmkRuleStub;
-import com.jetbrains.snakecharm.lang.psi.stubs.SmkSubworkflowStub;
+import com.jetbrains.snakecharm.lang.psi.*;
+import com.jetbrains.snakecharm.lang.psi.impl.stubs.*;
+import com.jetbrains.snakecharm.lang.psi.stubs.*;
 
 public interface SmkStubElementTypes {
     IStubElementType<SmkRuleStub, SmkRule>  RULE_DECLARATION_STATEMENT = new SmkRuleElementType();
     IStubElementType<SmkCheckpointStub, SmkCheckPoint> CHECKPOINT_DECLARATION_STATEMENT = new SmkCheckpointElementType();
     IStubElementType<SmkSubworkflowStub, SmkSubworkflow> SUBWORKFLOW_DECLARATION_STATEMENT = new SmkSubworkflowElementType();
+    IStubElementType<SmkModuleStub, SmkModule> MODULE_DECLARATION_STATEMENT = new SmkModuleElementType();
+    IStubElementType<SmkUseStub, SmkUse> USE_DECLARATION_STATEMENT = new SmkUseElementType();
 }

src/main/kotlin/com/jetbrains/snakecharm/codeInsight/SnakemakeAPI.kt

@@ -1,6 +1,12 @@
 package com.jetbrains.snakecharm.codeInsight
 
 import com.jetbrains.snakecharm.lang.SnakemakeNames
+import com.jetbrains.snakecharm.lang.SnakemakeNames.MODULE_CONFIG_KEYWORD
+import com.jetbrains.snakecharm.lang.SnakemakeNames.MODULE_META_WRAPPER_KEYWORD
+import com.jetbrains.snakecharm.lang.SnakemakeNames.MODULE_REPLACE_PREFIX_KEYWORD
+import com.jetbrains.snakecharm.lang.SnakemakeNames.MODULE_SKIP_VALIDATION_KEYWORD
+import com.jetbrains.snakecharm.lang.SnakemakeNames.MODULE_SNAKEFILE_KEYWORD
+import com.jetbrains.snakecharm.lang.SnakemakeNames.RULE_KEYWORD
 import com.jetbrains.snakecharm.lang.SnakemakeNames.SECTION_BENCHMARK
 import com.jetbrains.snakecharm.lang.SnakemakeNames.SECTION_CACHE
 import com.jetbrains.snakecharm.lang.SnakemakeNames.SECTION_CONDA
@@ -19,6 +25,7 @@ import com.jetbrains.snakecharm.lang.SnakemakeNames.SECTION_OUTPUT
 import com.jetbrains.snakecharm.lang.SnakemakeNames.SECTION_PARAMS
 import com.jetbrains.snakecharm.lang.SnakemakeNames.SECTION_PRIORITY
 import com.jetbrains.snakecharm.lang.SnakemakeNames.SECTION_RESOURCES
+import com.jetbrains.snakecharm.lang.SnakemakeNames.SECTION_RUN
 import com.jetbrains.snakecharm.lang.SnakemakeNames.SECTION_SCRIPT
 import com.jetbrains.snakecharm.lang.SnakemakeNames.SECTION_SHADOW
 import com.jetbrains.snakecharm.lang.SnakemakeNames.SECTION_SHELL
@@ -27,6 +34,9 @@ import com.jetbrains.snakecharm.lang.SnakemakeNames.SECTION_THREADS
 import com.jetbrains.snakecharm.lang.SnakemakeNames.SECTION_VERSION
 import com.jetbrains.snakecharm.lang.SnakemakeNames.SECTION_WILDCARD_CONSTRAINTS
 import com.jetbrains.snakecharm.lang.SnakemakeNames.SECTION_WRAPPER
+import com.jetbrains.snakecharm.lang.SnakemakeNames.SMK_AS_KEYWORD
+import com.jetbrains.snakecharm.lang.SnakemakeNames.SMK_FROM_KEYWORD
+import com.jetbrains.snakecharm.lang.SnakemakeNames.SMK_WITH_KEYWORD
 import com.jetbrains.snakecharm.lang.SnakemakeNames.SNAKEMAKE_IO_METHOD_ANCIENT
 import com.jetbrains.snakecharm.lang.SnakemakeNames.SNAKEMAKE_IO_METHOD_DIRECTORY
 import com.jetbrains.snakecharm.lang.SnakemakeNames.SNAKEMAKE_IO_METHOD_DYNAMIC
@@ -121,8 +131,8 @@ object SnakemakeAPI {
         SECTION_NAME,
         SECTION_HANDOVER
     )
-    val RULE_OR_CHECKPOINT_SECTION_KEYWORDS =
-        (RULE_OR_CHECKPOINT_ARGS_SECTION_KEYWORDS + setOf(SnakemakeNames.SECTION_RUN))
+
+    val RULE_OR_CHECKPOINT_SECTION_KEYWORDS = (RULE_OR_CHECKPOINT_ARGS_SECTION_KEYWORDS + setOf(SECTION_RUN))
 
     /**
      * For subworkflows parsing
@@ -133,6 +143,29 @@ object SnakemakeAPI {
         SnakemakeNames.SUBWORKFLOW_CONFIGFILE_KEYWORD
     )
 
+    /**
+     * For modules parsing
+     */
+    val MODULE_SECTIONS_KEYWORDS = setOf(
+        MODULE_SNAKEFILE_KEYWORD,
+        MODULE_CONFIG_KEYWORD,
+        MODULE_SKIP_VALIDATION_KEYWORD,
+        MODULE_META_WRAPPER_KEYWORD,
+        MODULE_REPLACE_PREFIX_KEYWORD
+    )
+
+    /**
+     * For uses parsing
+     */
+    val USE_SECTIONS_KEYWORDS = RULE_OR_CHECKPOINT_SECTION_KEYWORDS - EXECUTION_SECTIONS_KEYWORDS - SECTION_RUN
+
+    val USE_DECLARATION_KEYWORDS = setOf(
+        RULE_KEYWORD,
+        SMK_FROM_KEYWORD,
+        SMK_AS_KEYWORD,
+        SMK_WITH_KEYWORD
+    )
+
     /**
      * For type inference:
      * Some sections in snakemake are inaccessible after `rules.NAME.<section>`, so this set is required

src/main/kotlin/com/jetbrains/snakecharm/codeInsight/completion/SmkKeywordCompletionContributor.kt

@@ -9,7 +9,6 @@ import com.intellij.patterns.PlatformPatterns
 import com.intellij.patterns.PlatformPatterns.psiElement
 import com.intellij.patterns.PsiElementPattern
 import com.intellij.patterns.StandardPatterns
-import com.intellij.profile.codeInspection.InspectionProfileManager
 import com.intellij.psi.PsiComment
 import com.intellij.psi.PsiElement
 import com.intellij.psi.util.PsiTreeUtil
@@ -18,14 +17,19 @@ import com.intellij.util.ProcessingContext
 import com.jetbrains.python.PyTokenTypes
 import com.jetbrains.python.codeInsight.completion.PythonLookupElement
 import com.jetbrains.python.psi.*
+import com.jetbrains.snakecharm.codeInsight.SnakemakeAPI.MODULE_SECTIONS_KEYWORDS
 import com.jetbrains.snakecharm.codeInsight.SnakemakeAPI.RULE_OR_CHECKPOINT_SECTION_KEYWORDS
 import com.jetbrains.snakecharm.codeInsight.SnakemakeAPI.SUBWORKFLOW_SECTIONS_KEYWORDS
-import com.jetbrains.snakecharm.inspections.SmkUnrecognizedSectionInspection
+import com.jetbrains.snakecharm.codeInsight.SnakemakeAPI.USE_DECLARATION_KEYWORDS
+import com.jetbrains.snakecharm.codeInsight.SnakemakeAPI.USE_SECTIONS_KEYWORDS
 import com.jetbrains.snakecharm.lang.SnakemakeLanguageDialect
+import com.jetbrains.snakecharm.lang.SnakemakeNames
 import com.jetbrains.snakecharm.lang.parser.SmkTokenTypes.RULE_LIKE
+import com.jetbrains.snakecharm.lang.parser.SmkTokenTypes.USE_KEYWORD
 import com.jetbrains.snakecharm.lang.parser.SmkTokenTypes.WORKFLOW_TOPLEVEL_DECORATORS_WO_RULE_LIKE
 import com.jetbrains.snakecharm.lang.parser.SnakemakeLexer
 import com.jetbrains.snakecharm.lang.psi.*
+import com.jetbrains.snakecharm.lang.psi.impl.SmkUseArgsSectionImpl
 
 /**
  * @author Roman.Chernyatchik
@@ -54,6 +58,25 @@ class SmkKeywordCompletionContributor : CompletionContributor() {
             SubworkflowSectionKeywordsProvider.CAPTURE,
             SubworkflowSectionKeywordsProvider
         )
+
+        extend(
+            CompletionType.BASIC,
+            ModuleSectionKeywordsProvider.CAPTURE,
+            ModuleSectionKeywordsProvider
+        )
+
+        extend(
+            CompletionType.BASIC,
+            UseSectionKeywordsProvider.CAPTURE,
+            UseSectionKeywordsProvider
+        )
+    }
+}
+
+object RuleKeywordTail : TailType() {
+    override fun processTail(editor: Editor, tailOffset: Int): Int {
+        editor.document.insertString(tailOffset, " ${SnakemakeNames.RULE_KEYWORD} ")
+        return moveCaret(editor, tailOffset, 2 + SnakemakeNames.RULE_KEYWORD.length)
     }
 }
 
@@ -95,18 +118,19 @@ object WorkflowTopLevelKeywordsProvider : CompletionProvider<CompletionParameter
         val tokenType2Name = SnakemakeLexer.KEYWORDS
             .map { (k, v) -> v to k }
             .toMap()
-
         listOf(
             WORKFLOW_TOPLEVEL_DECORATORS_WO_RULE_LIKE to ColonAndWhiteSpaceTail,
             RULE_LIKE to TailType.SPACE
         ).forEach { (tokenSet, tail) ->
             tokenSet.types.forEach { tt ->
                 val s = tokenType2Name[tt]!!
 
+                val modifiedTail = if (tt == USE_KEYWORD) RuleKeywordTail else tail
+
                 result.addElement(
                     SmkCompletionUtil.createPrioritizedLookupElement(
                         TailTypeDecorator.withTail(
-                            PythonLookupElement(s, true, null), tail
+                            PythonLookupElement(s, true, null), modifiedTail
                         ),
                         SmkCompletionUtil.KEYWORDS_PRIORITY
                     )
@@ -152,7 +176,10 @@ object RuleSectionKeywordsProvider : CompletionProvider<CompletionParameters>()
         .inside(SmkRuleOrCheckpoint::class.java)!!
         .andNot(
             psiElement().insideOneOf(
-                PyArgumentList::class.java, SmkRunSection::class.java, PsiComment::class.java
+                PyArgumentList::class.java,
+                SmkRunSection::class.java,
+                PsiComment::class.java,
+                SmkUseArgsSectionImpl::class.java
             )
         )
 
@@ -204,6 +231,73 @@ object SubworkflowSectionKeywordsProvider : CompletionProvider<CompletionParamet
     }
 }
 
+object ModuleSectionKeywordsProvider : CompletionProvider<CompletionParameters>() {
+    val CAPTURE = psiElement()
+        .inFile(SmkKeywordCompletionContributor.IN_SNAKEMAKE)
+        .inside(psiElement().inside(SmkModule::class.java))
+        .andNot(
+            psiElement().insideOneOf(PyArgumentList::class.java, PsiComment::class.java)
+        )
+
+    override fun addCompletions(
+        parameters: CompletionParameters,
+        context: ProcessingContext,
+        result: CompletionResultSet
+    ) {
+        MODULE_SECTIONS_KEYWORDS.forEach { s ->
+            result.addElement(
+                SmkCompletionUtil.createPrioritizedLookupElement(
+                    TailTypeDecorator.withTail(
+                        PythonLookupElement(s, true, PlatformIcons.PROPERTY_ICON),
+                        ColonAndWhiteSpaceTail
+                    ),
+                    priority = SmkCompletionUtil.SECTIONS_KEYS_PRIORITY
+                )
+            )
+        }
+    }
+}
+
+object UseSectionKeywordsProvider : CompletionProvider<CompletionParameters>() {
+    val CAPTURE = psiElement()
+        .inFile(SmkKeywordCompletionContributor.IN_SNAKEMAKE)
+        .inside(psiElement().inside(SmkUse::class.java))
+        .andNot(
+            psiElement().insideOneOf(PyArgumentList::class.java, PsiComment::class.java)
+        )
+
+    override fun addCompletions(
+        parameters: CompletionParameters,
+        context: ProcessingContext,
+        result: CompletionResultSet
+    ) {
+        USE_DECLARATION_KEYWORDS.forEach { s ->
+            val tail = if (s != SnakemakeNames.SMK_WITH_KEYWORD) TailType.SPACE else TailType.CASE_COLON
+            result.addElement(
+                SmkCompletionUtil.createPrioritizedLookupElement(
+                    TailTypeDecorator.withTail(
+                        PythonLookupElement(s, true, PlatformIcons.PROPERTY_ICON),
+                        tail
+                    ),
+                    priority = SmkCompletionUtil.SECTIONS_KEYS_PRIORITY
+                )
+            )
+        }
+
+        USE_SECTIONS_KEYWORDS.forEach { s ->
+            result.addElement(
+                SmkCompletionUtil.createPrioritizedLookupElement(
+                    TailTypeDecorator.withTail(
+                        PythonLookupElement(s, true, PlatformIcons.PROPERTY_ICON),
+                        ColonAndWhiteSpaceTail
+                    ),
+                    priority = SmkCompletionUtil.SECTIONS_KEYS_PRIORITY
+                )
+            )
+        }
+    }
+}
+
 fun PsiElementPattern.Capture<PsiElement>.insideOneOf(
     vararg classes: Class<out PsiElement>
 ) = inside(

src/main/kotlin/com/jetbrains/snakecharm/inspections/SmkModuleRedeclarationInspection.kt

@@ -0,0 +1,32 @@
+package com.jetbrains.snakecharm.inspections
+
+import com.intellij.codeInspection.LocalInspectionToolSession
+import com.intellij.codeInspection.ProblemHighlightType
+import com.intellij.codeInspection.ProblemsHolder
+import com.jetbrains.snakecharm.SnakemakeBundle
+import com.jetbrains.snakecharm.lang.psi.SmkFile
+import com.jetbrains.snakecharm.lang.psi.SmkModule
+
+class SmkModuleRedeclarationInspection : SnakemakeInspection() {
+    override fun buildVisitor(
+        holder: ProblemsHolder,
+        isOnTheFly: Boolean,
+        session: LocalInspectionToolSession
+    ) = object : SnakemakeInspectionVisitor(holder, session) {
+        private val modulesNameAndPsi = (session.file as? SmkFile)?.collectModules() ?: emptyList()
+
+        override fun visitSmkModule(module: SmkModule) {
+            val name = module.name ?: return
+
+            if (module !== modulesNameAndPsi.findLast { it.first == name }?.second) {
+                registerProblem(
+                    module.originalElement,
+                    SnakemakeBundle.message("INSP.NAME.module.redeclaration"),
+                    ProblemHighlightType.LIKE_UNUSED_SYMBOL
+                )
+            }
+        }
+    }
+
+    override fun getDisplayName(): String = SnakemakeBundle.message("INSP.NAME.module.redeclaration")
+}

src/main/kotlin/com/jetbrains/snakecharm/inspections/SmkRuleRedeclarationInspection.kt

@@ -5,20 +5,22 @@ import com.intellij.codeInspection.ProblemHighlightType
 import com.intellij.codeInspection.ProblemHighlightType.GENERIC_ERROR
 import com.intellij.codeInspection.ProblemHighlightType.WEAK_WARNING
 import com.intellij.codeInspection.ProblemsHolder
+import com.intellij.psi.PsiElement
 import com.intellij.psi.util.PsiTreeUtil
 import com.jetbrains.python.psi.PyStatement
 import com.jetbrains.snakecharm.SnakemakeBundle
 import com.jetbrains.snakecharm.inspections.quickfix.RenameElementWithoutUsagesQuickFix
 import com.jetbrains.snakecharm.lang.psi.*
 import com.jetbrains.snakecharm.lang.psi.stubs.SmkCheckpointNameIndex
 import com.jetbrains.snakecharm.lang.psi.stubs.SmkRuleNameIndex
+import com.jetbrains.snakecharm.lang.psi.stubs.SmkUseNameIndex
 import com.jetbrains.snakecharm.lang.psi.types.AbstractSmkRuleOrCheckpointType
 
 class SmkRuleRedeclarationInspection : SnakemakeInspection() {
     override fun buildVisitor(
-            holder: ProblemsHolder,
-            isOnTheFly: Boolean,
-            session: LocalInspectionToolSession
+        holder: ProblemsHolder,
+        isOnTheFly: Boolean,
+        session: LocalInspectionToolSession
     ) = object : SnakemakeInspectionVisitor(holder, session) {
         private val localRules by lazy {
             holder.file.let { psiFile ->
@@ -40,6 +42,16 @@ class SmkRuleRedeclarationInspection : SnakemakeInspection() {
                 }
             }
         }
+        private val localUses by lazy {
+            holder.file.let { psiFile ->
+                when (psiFile) {
+                    is SmkFile -> {
+                        psiFile.collectUses().map { it.second }
+                    }
+                    else -> emptyList()
+                }
+            }
+        }
 
         override fun visitSmkRule(rule: SmkRule) {
             visitSMKRuleLike(rule)
@@ -49,6 +61,10 @@ class SmkRuleRedeclarationInspection : SnakemakeInspection() {
             visitSMKRuleLike(checkPoint)
         }
 
+        override fun visitSmkUse(use: SmkUse) {
+            visitSMKRuleLike(use)
+        }
+
         private fun visitSMKRuleLike(ruleLike: SmkRuleLike<SmkRuleOrCheckpointArgsSection>) {
             val containingFile = ruleLike.containingFile
             val nameToCheck = ruleLike.name ?: return
@@ -61,11 +77,17 @@ class SmkRuleRedeclarationInspection : SnakemakeInspection() {
                 ruleLike, nameToCheck, SmkCheckpointNameIndex.KEY, SmkCheckPoint::class.java
             ) { localCheckpoints }
 
-            val resolveResults = ruleResolveResults + cpResolveResults
+            val usesResolveResults: Collection<PsiElement> =
+                AbstractSmkRuleOrCheckpointType.findAvailableRuleLikeElementByName(
+                    ruleLike, nameToCheck, SmkUseNameIndex.KEY, SmkUse::class.java
+                ) { localUses }
+
+            val resolveResults = ruleResolveResults + cpResolveResults + usesResolveResults
 
             if (resolveResults.isEmpty()) {
                 return
             }
+
             if (resolveResults.size == 1 && resolveResults.single() == ruleLike) {
                 return
             }
@@ -98,7 +120,7 @@ class SmkRuleRedeclarationInspection : SnakemakeInspection() {
         }
 
         private fun registerProblem(
-            ruleLike: SmkRuleLike<SmkRuleOrCheckpointArgsSection>,
+            ruleLike: SmkSection,
             msg: String,
             severity: ProblemHighlightType
         ) {
@@ -114,5 +136,6 @@ class SmkRuleRedeclarationInspection : SnakemakeInspection() {
             )
         }
     }
+
     override fun getDisplayName(): String = SnakemakeBundle.message("INSP.NAME.rule.redeclaration")
 }