Support strings and tuple of strings in setup.py requirements (PY-29153, PY-30016)

python/src/com/jetbrains/python/packaging/PyPackageUtil.java

@@ -42,6 +42,7 @@ import com.jetbrains.python.PyNames;
 import com.jetbrains.python.codeInsight.controlflow.ScopeOwner;
 import com.jetbrains.python.packaging.setupPy.SetupTaskIntrospector;
 import com.jetbrains.python.psi.*;
+import com.jetbrains.python.psi.impl.PyPsiUtils;
 import com.jetbrains.python.psi.resolve.PyResolveContext;
 import com.jetbrains.python.psi.resolve.QualifiedResolveResult;
 import com.jetbrains.python.psi.types.TypeEvalContext;
@@ -171,21 +172,11 @@ public class PyPackageUtil {
   @Nullable
   private static Pair<String, List<PyRequirement>> getExtraRequires(@NotNull PyExpression extra, @Nullable PyExpression requires) {
     if (extra instanceof PyStringLiteralExpression) {
-      final String requiresValue;
-
-      if (requires instanceof PyStringLiteralExpression) {
-        requiresValue = ((PyStringLiteralExpression)requires).getStringValue();
-      }
-      else if (requires instanceof PyListLiteralExpression) {
-        final List<String> requiresListValue = PyUtil.strListValue(requires);
-        requiresValue = requiresListValue != null ? StringUtil.join(requiresListValue, "\n") : null;
-      }
-      else {
-        requiresValue = null;
-      }
+      final List<String> requiresValue = resolveRequiresValue(requires);
 
       if (requiresValue != null) {
-        return Pair.createNonNull(((PyStringLiteralExpression)extra).getStringValue(), PyRequirementParser.fromText(requiresValue));
+        return Pair.createNonNull(((PyStringLiteralExpression)extra).getStringValue(),
+                                  PyRequirementParser.fromText(StringUtil.join(requiresValue, "\n")));
       }
     }
 
@@ -199,11 +190,7 @@ public class PyPackageUtil {
       StreamEx
         .of(argumentNames)
         .map(setupCall::getKeywordArgument)
-        .map(requires -> resolveValue(requires, PyListLiteralExpression.class))
-        .nonNull()
-        .flatMap(requires -> Stream.of(requires.getElements()))
-        .select(PyStringLiteralExpression.class)
-        .map(StringLiteralExpression::getStringValue)
+        .flatCollection(PyPackageUtil::resolveRequiresValue)
         .joining("\n")
     );
   }
@@ -242,6 +229,32 @@ public class PyPackageUtil {
     return null;
   }
 
+  @Nullable
+  private static List<String> resolveRequiresValue(@Nullable PyExpression expression) {
+    PsiElement elementToAnalyze = PyPsiUtils.flattenParens(expression);
+
+    if (elementToAnalyze instanceof PyReferenceExpression) {
+      final TypeEvalContext context = TypeEvalContext.deepCodeInsight(elementToAnalyze.getProject());
+      final PyResolveContext resolveContext = PyResolveContext.noImplicits().withTypeEvalContext(context);
+      final QualifiedResolveResult result = ((PyReferenceExpression)elementToAnalyze).followAssignmentsChain(resolveContext);
+      elementToAnalyze = result.getElement();
+    }
+
+    if (elementToAnalyze instanceof PyStringLiteralExpression) {
+      return Collections.singletonList(((PyStringLiteralExpression)elementToAnalyze).getStringValue());
+    }
+    else if (elementToAnalyze instanceof PyListLiteralExpression || elementToAnalyze instanceof PyTupleExpression) {
+      return StreamEx
+        .of(((PySequenceExpression)elementToAnalyze).getElements())
+        .map(e -> resolveValue(e, PyStringLiteralExpression.class))
+        .select(PyStringLiteralExpression.class)
+        .map(PyStringLiteralExpression::getStringValue)
+        .toList();
+    }
+
+    return null;
+  }
+
   @NotNull
   public static List<String> getPackageNames(@NotNull Module module) {
     // TODO: Cache found module packages, clear cache on module updates

python/testData/packaging/PyPackageUtil/SetupPyExtrasReading/setup.py

@@ -0,0 +1,17 @@
+from distutils.core import setup
+
+tests_require = [
+    'mynose'
+]
+
+setup(name='foo',
+      version=0.1,
+      requires=[
+          'Markdown',
+      ],
+      install_requires='NewDjango==1.3.1',
+      tests_require=tests_require,
+      setup_requires=(
+          'numpy',
+      ),
+      extras_require={"e1": "r1", "e2": ("r2",), "e3": ["r3", "r4"]})

python/testData/packaging/PyPackageUtil/SetupPyReading/setup.py

@@ -9,10 +9,8 @@ setup(name='foo',
       requires=[
           'Markdown',
       ],
-      install_requires=[
-          'NewDjango==1.3.1',
-      ],
+      install_requires='NewDjango==1.3.1',
       tests_require=tests_require,
-      setup_requires=[
-          'numpy'
-      ])
+      setup_requires=(
+          'numpy',
+      ))

python/testSrc/com/jetbrains/python/packaging/PyPackageUtilTest.java

@@ -15,6 +15,7 @@
  */
 package com.jetbrains.python.packaging;
 
+import com.google.common.collect.ImmutableMap;
 import com.intellij.openapi.command.WriteCommandAction;
 import com.intellij.openapi.module.Module;
 import com.jetbrains.python.fixtures.PyTestCase;
@@ -25,7 +26,10 @@ import com.jetbrains.python.psi.PyFile;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
+import java.util.Arrays;
+import java.util.Collections;
 import java.util.List;
+import java.util.Map;
 
 public class PyPackageUtilTest extends PyTestCase {
 
@@ -37,19 +41,23 @@ public class PyPackageUtilTest extends PyTestCase {
   }
 
   public void testAbsentSetupPyReading() {
-    doTestSetupPyReading(false, false, false);
+    doTestSetupPyReading(false, false, false, false);
   }
 
   public void testAbsentSetupCallReading() {
-    doTestSetupPyReading(true, false, false);
+    doTestSetupPyReading(true, false, false, false);
   }
 
   public void testAbsentSetupPyRequiresReading() {
-    doTestSetupPyReading(true, true, false);
+    doTestSetupPyReading(true, true, false, false);
   }
 
   public void testSetupPyReading() {
-    doTestSetupPyReading(true, true, true);
+    doTestSetupPyReading(true, true, true, false);
+  }
+
+  public void testSetupPyExtrasReading() {
+    doTestSetupPyReading(true, true, true, true);
   }
 
   // PY-18966
@@ -104,7 +112,7 @@ public class PyPackageUtilTest extends PyTestCase {
     checkRequirements(PyPackageUtil.getRequirementsFromTxt(module));
   }
 
-  private void doTestSetupPyReading(boolean hasFile, boolean hasCall, boolean requires) {
+  private void doTestSetupPyReading(boolean hasFile, boolean hasCall, boolean requires, boolean extrasRequire) {
     final Module module = myFixture.getModule();
 
     if (hasFile) {
@@ -124,6 +132,21 @@ public class PyPackageUtilTest extends PyTestCase {
           assertNotNull(requirements);
           assertEmpty(requirements);
         }
+
+        if (extrasRequire) {
+          final Map<String, List<PyRequirement>> extrasRequirements = PyPackageUtil.findSetupPyExtrasRequire(module);
+
+          final ImmutableMap<String, List<PyRequirement>> expected = ImmutableMap.of(
+            "e1", Collections.singletonList(PyRequirementsKt.pyRequirement("r1")),
+            "e2", Collections.singletonList(PyRequirementsKt.pyRequirement("r2")),
+            "e3", Arrays.asList(PyRequirementsKt.pyRequirement("r3"), PyRequirementsKt.pyRequirement("r4"))
+          );
+
+          assertEquals(expected, extrasRequirements);
+        }
+        else {
+          assertNull(PyPackageUtil.findSetupPyExtrasRequire(module));
+        }
       }
       else {
         assertNull(PyPackageUtil.findSetupCall(module));