diff --git a/python/src/com/jetbrains/python/PythonFoldingBuilder.java b/python/src/com/jetbrains/python/PythonFoldingBuilder.java index 4127806be73b..ad3fc69879b4 100644 --- a/python/src/com/jetbrains/python/PythonFoldingBuilder.java +++ b/python/src/com/jetbrains/python/PythonFoldingBuilder.java @@ -26,8 +26,10 @@ import com.intellij.openapi.util.text.LineTokenizer; import com.intellij.openapi.util.text.StringUtil; import com.intellij.psi.PsiElement; import com.intellij.psi.tree.IElementType; +import com.intellij.psi.tree.TokenSet; import com.jetbrains.python.psi.*; import com.jetbrains.python.psi.impl.PyFileImpl; +import com.jetbrains.python.psi.impl.PyStringLiteralExpressionImpl; import org.jetbrains.annotations.NotNull; import org.jetbrains.annotations.Nullable; @@ -38,6 +40,16 @@ import java.util.List; */ public class PythonFoldingBuilder extends CustomFoldingBuilder implements DumbAware { + public static final TokenSet FOLDABLE_COLLECTIONS_LITERALS = TokenSet.create( + PyElementTypes.SET_LITERAL_EXPRESSION, + PyElementTypes.DICT_LITERAL_EXPRESSION, + PyElementTypes.GENERATOR_EXPRESSION, + PyElementTypes.SET_COMP_EXPRESSION, + PyElementTypes.DICT_COMP_EXPRESSION, + PyElementTypes.LIST_LITERAL_EXPRESSION, + PyElementTypes.LIST_COMP_EXPRESSION, + PyElementTypes.TUPLE_EXPRESSION); + @Override protected void buildLanguageFoldRegions(@NotNull List descriptors, @NotNull PsiElement root, @@ -47,7 +59,8 @@ public class PythonFoldingBuilder extends CustomFoldingBuilder implements DumbAw } private static void appendDescriptors(ASTNode node, List descriptors) { - if (node.getElementType() instanceof PyFileElementType) { + IElementType elementType = node.getElementType(); + if (elementType instanceof PyFileElementType) { final List imports = ((PyFile)node.getPsi()).getImportBlock(); if (imports.size() > 1) { final PyImportStatementBase firstImport = imports.get(0); @@ -56,13 +69,15 @@ public class PythonFoldingBuilder extends CustomFoldingBuilder implements DumbAw lastImport.getTextRange().getEndOffset()))); } } - else if (node.getElementType() == PyElementTypes.STATEMENT_LIST) { + else if (elementType == PyElementTypes.STATEMENT_LIST) { foldStatementList(node, descriptors); } - else if (node.getElementType() == PyElementTypes.STRING_LITERAL_EXPRESSION) { - foldDocString(node, descriptors); + else if (elementType == PyElementTypes.STRING_LITERAL_EXPRESSION) { + foldLongStrings(node, descriptors); + } + else if (FOLDABLE_COLLECTIONS_LITERALS.contains(elementType)) { + foldCollectionLiteral(node, descriptors); } - ASTNode child = node.getFirstChildNode(); while (child != null) { appendDescriptors(child, descriptors); @@ -70,6 +85,14 @@ public class PythonFoldingBuilder extends CustomFoldingBuilder implements DumbAw } } + private static void foldCollectionLiteral(ASTNode node, List descriptors) { + if (StringUtil.countNewLines(node.getChars()) > 0) { + TextRange range = node.getTextRange(); + int delta = node.getElementType() == PyElementTypes.TUPLE_EXPRESSION ? 0 : 1; + descriptors.add(new FoldingDescriptor(node, TextRange.create(range.getStartOffset() + delta, range.getEndOffset() - delta))); + } + } + private static void foldStatementList(ASTNode node, List descriptors) { IElementType elType = node.getTreeParent().getElementType(); if (elType == PyElementTypes.FUNCTION_DECLARATION @@ -106,8 +129,11 @@ public class PythonFoldingBuilder extends CustomFoldingBuilder implements DumbAw return false; } - private static void foldDocString(ASTNode node, List descriptors) { - if (getDocStringOwnerType(node) != null && StringUtil.countChars(node.getText(), '\n') > 1) { + private static void foldLongStrings(ASTNode node, List descriptors) { + //don't want to fold docstrings like """\n string \n """ + boolean shouldFoldDocString = getDocStringOwnerType(node) != null && StringUtil.countNewLines(node.getChars()) > 1; + boolean shouldFoldString = getDocStringOwnerType(node) == null && StringUtil.countNewLines(node.getChars()) > 0; + if (shouldFoldDocString || shouldFoldString) { descriptors.add(new FoldingDescriptor(node, node.getTextRange())); } } @@ -137,12 +163,33 @@ public class PythonFoldingBuilder extends CustomFoldingBuilder implements DumbAw return "import ..."; } if (node.getElementType() == PyElementTypes.STRING_LITERAL_EXPRESSION) { - final String stringValue = ((PyStringLiteralExpression)node.getPsi()).getStringValue().trim(); - final String[] lines = LineTokenizer.tokenize(stringValue, true); - if (lines.length > 2 && lines[1].trim().length() == 0) { - return "\"\"\"" + lines [0].trim() + "...\"\"\""; + PyStringLiteralExpression stringLiteralExpression = (PyStringLiteralExpression)node.getPsi(); + if (stringLiteralExpression.isDocString()) { + final String stringValue = stringLiteralExpression.getStringValue().trim(); + final String[] lines = LineTokenizer.tokenize(stringValue, true); + if (lines.length > 2 && lines[1].trim().length() == 0) { + return "\"\"\"" + lines[0].trim() + "...\"\"\""; + } + return "\"\"\"...\"\"\""; + } else { + return getLanguagePlaceholderForString(stringLiteralExpression); } - return "\"\"\"...\"\"\""; + } + return "..."; + } + + private static String getLanguagePlaceholderForString(PyStringLiteralExpression stringLiteralExpression) { + String stringText = stringLiteralExpression.getText(); + int prefixLength = PyStringLiteralExpressionImpl.getPrefixLength(stringText); + stringText = stringText.substring(prefixLength); + if (stringText.startsWith("'''") && stringText.endsWith("'''")) { + return "'''...'''"; + } + if (stringText.startsWith("'") && stringText.endsWith("'")) { + return "'...'"; + } + if (stringText.startsWith("\"") && stringText.endsWith("\"")) { + return "\"...\""; } return "..."; } diff --git a/python/testData/folding/collectionsFolding.py b/python/testData/folding/collectionsFolding.py new file mode 100644 index 000000000000..0ec4ffbafb0e --- /dev/null +++ b/python/testData/folding/collectionsFolding.py @@ -0,0 +1,9 @@ +pyt = ((x, y, z) for z in b + for y in xrange(1, z) + for x in range(1, y) if x * x + y * y == z * z) + +primes = {x for x in range(2, 101) + if all(x % y for y in range(2, min(x, 11)))} + +d = {n: n ** 2 + for n in range(5)} \ No newline at end of file diff --git a/python/testData/folding/longStringsFolding.py b/python/testData/folding/longStringsFolding.py new file mode 100644 index 000000000000..96ae0bfd38c8 --- /dev/null +++ b/python/testData/folding/longStringsFolding.py @@ -0,0 +1,4 @@ +print("This is " + "really " + "long string. " + "It should be foldable") \ No newline at end of file diff --git a/python/testSrc/com/jetbrains/python/PyFoldingTest.java b/python/testSrc/com/jetbrains/python/PyFoldingTest.java index 201376efcd9f..a48c7c032adf 100644 --- a/python/testSrc/com/jetbrains/python/PyFoldingTest.java +++ b/python/testSrc/com/jetbrains/python/PyFoldingTest.java @@ -44,4 +44,13 @@ public class PyFoldingTest extends PyTestCase { public void testBlocksFolding() { doTest(); } + + public void testLongStringsFolding() { + doTest(); + } + + public void testCollectionsFolding() { + doTest(); + } + }