From 714f57764d365ad50cb1372203f8c5a65338c604 Mon Sep 17 00:00:00 2001
From: Morgan Bartholomew <morgan.bartholomew@jetbrains.com>
Date: Sun, 2 Nov 2025 13:02:12 +1000
Subject: [PATCH] [python] PY-21499 support regex injection for `regex` module

(cherry picked from commit 3f5a4db6967534e746650b8ea10ad4d6a29a0363)

GitOrigin-RevId: 4db040ead1a80f49b78378267aeb35b350818069
---
 .../regexp/PythonRegexpInjector.java          |  8 ++-
 .../com/jetbrains/python/PyRegexTest.kt       | 50 +++++++++++++++++++
 2 files changed, 56 insertions(+), 2 deletions(-)
 create mode 100644 python/testSrc/com/jetbrains/python/PyRegexTest.kt

diff --git a/python/src/com/jetbrains/python/codeInsight/regexp/PythonRegexpInjector.java b/python/src/com/jetbrains/python/codeInsight/regexp/PythonRegexpInjector.java
index e8e2be33d619..9d6957cf82f5 100644
--- a/python/src/com/jetbrains/python/codeInsight/regexp/PythonRegexpInjector.java
+++ b/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;
     }
diff --git a/python/testSrc/com/jetbrains/python/PyRegexTest.kt b/python/testSrc/com/jetbrains/python/PyRegexTest.kt
new file mode 100644
index 000000000000..b62f996eeb3d
--- /dev/null
+++ b/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"
+        )
+      }
+    }
+  }
+}