diff --git a/python/src/com/jetbrains/python/psi/PyUtil.java b/python/src/com/jetbrains/python/psi/PyUtil.java index c4fe3563ae36..c7ca1e99b990 100644 --- a/python/src/com/jetbrains/python/psi/PyUtil.java +++ b/python/src/com/jetbrains/python/psi/PyUtil.java @@ -675,7 +675,11 @@ public class PyUtil { @Nullable public static List getStringListFromTargetExpression(PyTargetExpression attr) { - PyExpression value = attr.findAssignedValue(); + return strListValue(attr.findAssignedValue()); + } + + @Nullable + public static List strListValue(PyExpression value) { while (value instanceof PyParenthesizedExpression) { value = ((PyParenthesizedExpression) value).getContainedExpression(); } diff --git a/python/src/com/jetbrains/python/psi/impl/PyFileImpl.java b/python/src/com/jetbrains/python/psi/impl/PyFileImpl.java index a7c755211975..51de0c664174 100644 --- a/python/src/com/jetbrains/python/psi/impl/PyFileImpl.java +++ b/python/src/com/jetbrains/python/psi/impl/PyFileImpl.java @@ -33,10 +33,7 @@ import org.jetbrains.annotations.Nullable; import javax.swing.*; import java.lang.ref.SoftReference; -import java.util.ArrayList; -import java.util.HashSet; -import java.util.List; -import java.util.Set; +import java.util.*; public class PyFileImpl extends PsiFileBase implements PyFile, PyExpression { protected PyType myType; @@ -416,12 +413,45 @@ public class PyFileImpl extends PsiFileBase implements PyFile, PyExpression { private static class DunderAllBuilder extends PyRecursiveElementVisitor { private List myResult = null; private boolean myDynamic = false; + private boolean myFoundDunderAll = false; + + // hashlib builds __all__ by concatenating multiple lists of strings, and we want to understand this + private Map> myDunderLike = new HashMap>(); @Override public void visitPyTargetExpression(PyTargetExpression node) { if (PyNames.ALL.equals(node.getName())) { - myResult = PyUtil.getStringListFromTargetExpression(node); + myFoundDunderAll = true; + PyExpression value = node.findAssignedValue(); + if (value instanceof PyBinaryExpression) { + PyBinaryExpression binaryExpression = (PyBinaryExpression)value; + if (binaryExpression.isOperator("+")) { + List lhs = getStringListFromValue(binaryExpression.getLeftExpression()); + List rhs = getStringListFromValue(binaryExpression.getRightExpression()); + if (lhs != null && rhs != null) { + myResult = new ArrayList(lhs); + myResult.addAll(rhs); + } + } + } + else { + myResult = PyUtil.getStringListFromTargetExpression(node); + } } + if (!myFoundDunderAll) { + List names = PyUtil.getStringListFromTargetExpression(node); + if (names != null) { + myDunderLike.put(node.getName(), names); + } + } + } + + @Nullable + private List getStringListFromValue(PyExpression expression) { + if (expression instanceof PyReferenceExpression && ((PyReferenceExpression)expression).getQualifier() == null) { + return myDunderLike.get(((PyReferenceExpression)expression).getReferencedName()); + } + return PyUtil.strListValue(expression); } @Override diff --git a/python/testData/stubs/DunderAllAsSum.py b/python/testData/stubs/DunderAllAsSum.py new file mode 100644 index 000000000000..898ca7e32029 --- /dev/null +++ b/python/testData/stubs/DunderAllAsSum.py @@ -0,0 +1,4 @@ +__always_supported = ('md5', 'sha1') + +__all__ = __always_supported + ('algorithms_guaranteed', + 'algorithms_available') diff --git a/python/testSrc/com/jetbrains/python/PyStubsTest.java b/python/testSrc/com/jetbrains/python/PyStubsTest.java index b192f2b36592..cc4e961eb4aa 100644 --- a/python/testSrc/com/jetbrains/python/PyStubsTest.java +++ b/python/testSrc/com/jetbrains/python/PyStubsTest.java @@ -251,6 +251,13 @@ public class PyStubsTest extends PyLightFixtureTestCase { assertNotParsed(file); } + public void testDunderAllAsSum() { + final PyFileImpl file = (PyFileImpl) getTestFile(); + final List all = file.getDunderAll(); + assertSameElements(all, "md5", "sha1", "algorithms_guaranteed", "algorithms_available"); + assertNotParsed(file); + } + public void testSlots() { final PyFileImpl file = (PyFileImpl) getTestFile(); final PyClass pyClass = file.getTopLevelClasses().get(0);