From e5166b993e9ae6c433ec704415172eaeca78494c Mon Sep 17 00:00:00 2001
From: "lada.gagina" <lada.gagina@jetbrains.com>
Date: Fri, 5 Jan 2024 23:16:25 +0100
Subject: [PATCH] [python] Inject requirements language into pyproject.toml
 (PY-42296, PY-54440)

GitOrigin-RevId: 196f56270d0a102a074872871898d640d77b2918
---
 python/src/META-INF/python-core-common.xml    |  5 +--
 .../RequirementsLanguageInjector.kt           | 34 ++++++++++++++++++
 .../pyproject.toml                            |  2 ++
 .../pyproject.toml                            |  2 ++
 .../pyproject.toml                            |  2 ++
 .../pyproject.toml                            |  2 ++
 .../RequirementsCompletionTest.kt             | 36 +++++++++++++++++++
 7 files changed, 79 insertions(+), 4 deletions(-)
 create mode 100644 python/src/com/intellij/python/community/impl/requirements/RequirementsLanguageInjector.kt
 create mode 100644 python/testData/requirements/completion/packageNameInBuildSystemTablePyProjectToml/pyproject.toml
 create mode 100644 python/testData/requirements/completion/packageNameInProjectTablePyProjectToml/pyproject.toml
 create mode 100644 python/testData/requirements/completion/versionInBuildSystemTablePyProjectToml/pyproject.toml
 create mode 100644 python/testData/requirements/completion/versionInProjectTablePyProjectToml/pyproject.toml

diff --git a/python/src/META-INF/python-core-common.xml b/python/src/META-INF/python-core-common.xml
index cfb42eac4dcd..2b07a429103b 100644
--- a/python/src/META-INF/python-core-common.xml
+++ b/python/src/META-INF/python-core-common.xml
@@ -28,6 +28,7 @@
     <lang.syntaxHighlighterFactory
       language="Requirements"
       implementationClass="com.intellij.python.community.impl.requirements.RequirementsSyntaxHighlighterFactory"/>
+    <multiHostInjector implementation="com.intellij.python.community.impl.requirements.RequirementsLanguageInjector"/>
     <completion.contributor language="Requirements"
                             implementationClass="com.intellij.python.community.impl.requirements.RequirementsPackageNameCompletionContributor"/>
     <completion.contributor language="Requirements"
@@ -542,10 +543,6 @@
                        key="python.sdk.installation.notification.group"/>
 
 
-    <completion.contributor language="TOML"
-                            implementationClass="com.jetbrains.python.packaging.pyproject.PyprojectPackageCompletionContributor"/>
-    <completion.contributor language="TOML"
-                            implementationClass="com.jetbrains.python.packaging.pyproject.PyprojectMetadataCompletionContributor"/>
     <localInspection language="TOML"
                      shortName="PyprojectInspection"
                      suppressId="PyprojectInspection"
diff --git a/python/src/com/intellij/python/community/impl/requirements/RequirementsLanguageInjector.kt b/python/src/com/intellij/python/community/impl/requirements/RequirementsLanguageInjector.kt
new file mode 100644
index 000000000000..459b566f5fe9
--- /dev/null
+++ b/python/src/com/intellij/python/community/impl/requirements/RequirementsLanguageInjector.kt
@@ -0,0 +1,34 @@
+// Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.intellij.python.community.impl.requirements
+
+import com.intellij.lang.injection.MultiHostInjector
+import com.intellij.lang.injection.MultiHostRegistrar
+import com.intellij.openapi.util.TextRange
+import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiLanguageInjectionHost
+import org.toml.lang.psi.TomlKeyValue
+import org.toml.lang.psi.TomlLiteral
+import org.toml.lang.psi.TomlTable
+
+class RequirementsLanguageInjector : MultiHostInjector {
+  override fun getLanguagesToInject(registrar: MultiHostRegistrar, context: PsiElement) {
+    val parent = (context as TomlLiteral).parent.parent
+    if (parent is TomlKeyValue) {
+      if (parent.key.text == "dependencies" && "project" == (parent.parent as? TomlTable)?.header?.key?.text
+          || parent.key.text == "requires" && "build-system" == (parent.parent as? TomlTable)?.header?.key?.text) {
+        registrar.startInjecting(RequirementsLanguage.INSTANCE).addPlace(null, null, context as PsiLanguageInjectionHost,
+                                                                         TextRange.create(1,
+                                                                                          (context as PsiLanguageInjectionHost).textLength - 1)).doneInjecting()
+      }
+    }
+  }
+
+  override fun elementsToInjectIn(): MutableList<out Class<out PsiElement>> {
+    try {
+      return mutableListOf(TomlLiteral::class.java)
+    }
+    catch (exception: NoClassDefFoundError) {
+      return mutableListOf()
+    }
+  }
+}
\ No newline at end of file
diff --git a/python/testData/requirements/completion/packageNameInBuildSystemTablePyProjectToml/pyproject.toml b/python/testData/requirements/completion/packageNameInBuildSystemTablePyProjectToml/pyproject.toml
new file mode 100644
index 000000000000..0f00fe3516ce
--- /dev/null
+++ b/python/testData/requirements/completion/packageNameInBuildSystemTablePyProjectToml/pyproject.toml
@@ -0,0 +1,2 @@
+[build-system]
+requires = [ "my<caret>" ]
\ No newline at end of file
diff --git a/python/testData/requirements/completion/packageNameInProjectTablePyProjectToml/pyproject.toml b/python/testData/requirements/completion/packageNameInProjectTablePyProjectToml/pyproject.toml
new file mode 100644
index 000000000000..25ff840b737c
--- /dev/null
+++ b/python/testData/requirements/completion/packageNameInProjectTablePyProjectToml/pyproject.toml
@@ -0,0 +1,2 @@
+[project]
+dependencies = [ "my<caret>" ]
\ No newline at end of file
diff --git a/python/testData/requirements/completion/versionInBuildSystemTablePyProjectToml/pyproject.toml b/python/testData/requirements/completion/versionInBuildSystemTablePyProjectToml/pyproject.toml
new file mode 100644
index 000000000000..2bc287004656
--- /dev/null
+++ b/python/testData/requirements/completion/versionInBuildSystemTablePyProjectToml/pyproject.toml
@@ -0,0 +1,2 @@
+[build-system]
+requires = [ "mypy==<caret>" ]
\ No newline at end of file
diff --git a/python/testData/requirements/completion/versionInProjectTablePyProjectToml/pyproject.toml b/python/testData/requirements/completion/versionInProjectTablePyProjectToml/pyproject.toml
new file mode 100644
index 000000000000..25b6183f050c
--- /dev/null
+++ b/python/testData/requirements/completion/versionInProjectTablePyProjectToml/pyproject.toml
@@ -0,0 +1,2 @@
+[project]
+dependencies = [ "mypy==<caret>" ]
\ No newline at end of file
diff --git a/python/testSrc/com/intellij/python/community/impl/requirements/RequirementsCompletionTest.kt b/python/testSrc/com/intellij/python/community/impl/requirements/RequirementsCompletionTest.kt
index 8393c64b6918..7bed1d6ad468 100644
--- a/python/testSrc/com/intellij/python/community/impl/requirements/RequirementsCompletionTest.kt
+++ b/python/testSrc/com/intellij/python/community/impl/requirements/RequirementsCompletionTest.kt
@@ -30,6 +30,24 @@ class RequirementsCompletionTest : BasePlatformTestCase() {
     assertContainsElements(myFixture.lookupElementStrings!!, "mypy", "mypy-extensions")
   }
 
+  fun testPackageNameInProjectTablePyProjectToml() {
+    mockPackageNames(listOf("mypy", "mypy-extensions"))
+
+    myFixture.configureByFile("${getTestName(true)}/pyproject.toml")
+    myFixture.completeBasic()
+    assertNotEmpty(myFixture.lookupElementStrings)
+    assertContainsElements(myFixture.lookupElementStrings!!, "mypy", "mypy-extensions")
+  }
+
+  fun testPackageNameInBuildSystemTablePyProjectToml() {
+    mockPackageNames(listOf("mypy", "mypy-extensions"))
+
+    myFixture.configureByFile("${getTestName(true)}/pyproject.toml")
+    myFixture.completeBasic()
+    assertNotEmpty(myFixture.lookupElementStrings)
+    assertContainsElements(myFixture.lookupElementStrings!!, "mypy", "mypy-extensions")
+  }
+
   fun testVersionCompletion() {
     mockPackageDetails("mypy", listOf("1.7.0", "1.6.1", "1.6.0"))
 
@@ -39,6 +57,24 @@ class RequirementsCompletionTest : BasePlatformTestCase() {
     assertContainsOrdered(myFixture.lookupElementStrings!!, "1.7.0", "1.6.1", "1.6.0")
   }
 
+  fun testVersionInBuildSystemTablePyProjectToml() {
+    mockPackageDetails("mypy", listOf("1.7.0", "1.6.1", "1.6.0"))
+
+    myFixture.configureByFile("${getTestName(true)}/pyproject.toml")
+    myFixture.completeBasic()
+    assertNotEmpty(myFixture.lookupElementStrings)
+    assertContainsOrdered(myFixture.lookupElementStrings!!, "1.7.0", "1.6.1", "1.6.0")
+  }
+
+  fun testVersionInProjectTablePyProjectToml() {
+    mockPackageDetails("mypy", listOf("1.7.0", "1.6.1", "1.6.0"))
+
+    myFixture.configureByFile("${getTestName(true)}/pyproject.toml")
+    myFixture.completeBasic()
+    assertNotEmpty(myFixture.lookupElementStrings)
+    assertContainsOrdered(myFixture.lookupElementStrings!!, "1.7.0", "1.6.1", "1.6.0")
+  }
+
   private fun mockPackageNames(packageNames: List<String>) {
     val packageManagerProvider = TestPackageManagerProvider().withPackageNames(packageNames)
     ExtensionTestUtil.maskExtensions(PythonPackageManagerProvider.EP_NAME, listOf(packageManagerProvider), testRootDisposable)