[python] PY-21499 support regex injection for regex module

GitOrigin-RevId: 3f5a4db6967534e746650b8ea10ad4d6a29a0363

python/src/com/jetbrains/python/codeInsight/regexp/PythonRegexpInjector.java

@@ -37,12 +37,16 @@ public final class PythonRegexpInjector implements MultiHostInjector {
     addMethod("compile");
     addMethod("search");
     addMethod("match");
+    addMethod("fullmatch");
     addMethod("split");
+    addMethod("splititer");
     addMethod("findall");
     addMethod("finditer");
     addMethod("sub");
+    addMethod("subf");
     addMethod("subn");
-    addMethod("fullmatch");
+    addMethod("subfn");
+    addMethod("template");
   }
 
   private void addMethod(@NotNull String name) {
@@ -107,7 +111,7 @@ public final class PythonRegexpInjector implements MultiHostInjector {
   private @Nullable RegexpMethodDescriptor findRegexpMethodDescriptor(@Nullable PsiElement element) {
     if (element == null ||
         !(ScopeUtil.getScopeOwner(element) instanceof PyFile) ||
-        !ArrayUtil.contains(element.getContainingFile().getName(), "re.py", "re.pyi") ||
+        !ArrayUtil.contains(element.getContainingFile().getName(), "re.py", "re.pyi", "regex.py", "regex.pyi") ||
         !(element instanceof PyFunction)) {
       return null;
     }

python/testSrc/com/jetbrains/python/PyRegexTest.kt

@@ -0,0 +1,50 @@
+// Copyright 2000-2022 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.jetbrains.python
+
+import com.intellij.idea.TestFor
+import com.intellij.lang.injection.InjectedLanguageManager
+import com.intellij.psi.PsiElement
+import com.jetbrains.python.fixtures.PyTestCase
+import junit.framework.TestCase
+import org.junit.jupiter.params.ParameterizedTest
+import org.junit.jupiter.params.provider.ValueSource
+
+/**
+ *  tests for the `regex` package
+ */
+class PyRegexTest : PyTestCase() {
+  fun `test avoid class failing because there are no tests`() {
+    // TODO: delete this when we refactor to JUnit5
+  }
+
+  private fun doTestInjectedText(text: String, expected: String): PsiElement {
+    myFixture.configureByText(PythonFileType.INSTANCE, text)
+    val languageManager = InjectedLanguageManager.getInstance(myFixture.project)
+    val host = languageManager.getInjectionHost(elementAtCaret)
+    assertNotNull(host)
+    val files = languageManager.getInjectedPsiFiles(host!!)
+    assertNotNull(files)
+    assertFalse(files!!.isEmpty())
+    val injected = files[0]!!.first
+    TestCase.assertEquals(expected, injected.text)
+    return injected
+  }
+
+  @ParameterizedTest
+  @ValueSource(strings = ["compile", "splititer", "subf", "subfn", "template"])
+  @TestFor(issues = ["PY-21499"])
+  fun `language injection`(regexFunction: String) {
+    runBare {
+      runWithAdditionalFileInLibDir("regex.py", "def $regexFunction(): ...") {
+        doTestInjectedText(
+          """
+          import regex
+         
+          regex.$regexFunction("<caret>.*a")
+          """.trimIndent(),
+          ".*a"
+        )
+      }
+    }
+  }
+}