PY-31689 Support wrapping options for list literals

GitOrigin-RevId: be976ddc224983174f5566e53985e29fd49051f4

python/python-common-tests/com/jetbrains/python/PythonCommonFormatterTest.java

@@ -6,6 +6,7 @@ import com.intellij.openapi.command.WriteCommandAction;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFile;
 import com.intellij.psi.codeStyle.CodeStyleManager;
+import com.intellij.psi.codeStyle.CommonCodeStyleSettings;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.jetbrains.python.codeInsight.PyCodeInsightSettings;
 import com.jetbrains.python.fixture.PythonCommonTestCase;
@@ -1362,4 +1363,61 @@ public abstract class PythonCommonFormatterTest extends PythonCommonTestCase {
   public void testSpaceAfterStarInTypeParameterList() {
     doTest();
   }
+
+  // PY-31689
+  public void testListWrappingBracketsOnNewLine() {
+    getCodeStyleSettings().setRightMargin(PythonLanguage.getInstance(), 40);
+    getPythonCodeStyleSettings().LIST_WRAPPING = CommonCodeStyleSettings.WRAP_AS_NEEDED;
+    getPythonCodeStyleSettings().LIST_NEW_LINE_AFTER_LEFT_BRACKET = true;
+    getPythonCodeStyleSettings().LIST_NEW_LINE_BEFORE_RIGHT_BRACKET = true;
+    doTest();
+  }
+
+  // PY-31689
+  public void testListChopDownBracketsOnNewLine() {
+    getCodeStyleSettings().setRightMargin(PythonLanguage.getInstance(), 40);
+    getPythonCodeStyleSettings().LIST_WRAPPING = CommonCodeStyleSettings.WRAP_ON_EVERY_ITEM;
+    getPythonCodeStyleSettings().LIST_NEW_LINE_AFTER_LEFT_BRACKET = true;
+    getPythonCodeStyleSettings().LIST_NEW_LINE_BEFORE_RIGHT_BRACKET = true;
+    doTest();
+  }
+
+  // PY-31689
+  public void testListWrappingBracketsOnSameLine() {
+    getCodeStyleSettings().setRightMargin(PythonLanguage.getInstance(), 40);
+    getPythonCodeStyleSettings().LIST_WRAPPING = CommonCodeStyleSettings.WRAP_AS_NEEDED;
+    doTest();
+  }
+
+  // PY-31689
+  public void testListChopDownBracketsOnSameLine() {
+    getPythonCodeStyleSettings().LIST_WRAPPING = CommonCodeStyleSettings.WRAP_ON_EVERY_ITEM;
+    getCodeStyleSettings().setRightMargin(PythonLanguage.getInstance(), 40);
+    doTest();
+  }
+
+  // PY-31689
+  public void testListWrappingFirstBracketOnNewLine() {
+    getCodeStyleSettings().setRightMargin(PythonLanguage.getInstance(), 40);
+    getPythonCodeStyleSettings().LIST_WRAPPING = CommonCodeStyleSettings.WRAP_AS_NEEDED;
+    getPythonCodeStyleSettings().LIST_NEW_LINE_AFTER_LEFT_BRACKET = true;
+    doTest();
+  }
+
+  // PY-31689
+  public void testListWrappingLastBracketOnNewLine() {
+    getCodeStyleSettings().setRightMargin(PythonLanguage.getInstance(), 40);
+    getPythonCodeStyleSettings().LIST_WRAPPING = CommonCodeStyleSettings.WRAP_AS_NEEDED;
+    getPythonCodeStyleSettings().LIST_NEW_LINE_BEFORE_RIGHT_BRACKET = true;
+    doTest();
+  }
+
+  // PY-31689
+  public void testListWrappingDoNotAlign() {
+    getCodeStyleSettings().setRightMargin(PythonLanguage.getInstance(), 40);
+    getPythonCodeStyleSettings().LIST_WRAPPING = CommonCodeStyleSettings.WRAP_ON_EVERY_ITEM;
+    getPythonCodeStyleSettings().LIST_NEW_LINE_BEFORE_RIGHT_BRACKET = true;
+    getPythonCodeStyleSettings().ALIGN_COLLECTIONS_AND_COMPREHENSIONS = false;
+    doTest();
+  }
 }

python/python-syntax-core/src/com/jetbrains/python/formatter/PyBlock.java

@@ -96,6 +96,7 @@ public class PyBlock implements ASTBlock {
   private Alignment myChildAlignment = null;
   private Alignment myDictAlignment = null;
   private Wrap myDictWrapping = null;
+  private Wrap myListWrapping = null;
   private Wrap myFromImportWrapping = null;
   private Wrap myParameterListWrapping = null;
   private Wrap myArgumentListWrapping = null;
@@ -120,6 +121,9 @@ public class PyBlock implements ASTBlock {
       myDictAlignment = Alignment.createAlignment(true);
       myDictWrapping = Wrap.createWrap(pySettings.DICT_WRAPPING, true);
     }
+    else if (node.getElementType() == PyElementTypes.LIST_LITERAL_EXPRESSION) {
+      myListWrapping = Wrap.createWrap(pySettings.LIST_WRAPPING, pySettings.LIST_NEW_LINE_AFTER_LEFT_BRACKET);
+    }
     else if (node.getElementType() == PyElementTypes.FROM_IMPORT_STATEMENT) {
       myFromImportWrapping = Wrap.createWrap(pySettings.FROM_IMPORT_WRAPPING, false);
     }
@@ -428,6 +432,12 @@ public class PyBlock implements ASTBlock {
     if (childType == PyElementTypes.KEY_VALUE_EXPRESSION && isChildOfDictLiteral(child)) {
       childWrap = myDictWrapping;
     }
+    if (parentType == PyElementTypes.LIST_LITERAL_EXPRESSION &&
+        childType != PyTokenTypes.COMMA &&
+        childType != PyTokenTypes.LBRACKET &&
+        childType != PyTokenTypes.RBRACKET) {
+      childWrap = myListWrapping;
+    }
     if (parentType == PyElementTypes.PARAMETER_LIST &&
         childType != PyTokenTypes.COMMA &&
         childType != PyTokenTypes.LPAR &&
@@ -744,8 +754,11 @@ public class PyBlock implements ASTBlock {
           myContext.getMode() == FormattingMode.ADJUST_INDENT) {
         return true;
       }
-      return !hasHangingIndent(myNode.getPsi()) && !(myNode.getElementType() == PyElementTypes.DICT_LITERAL_EXPRESSION &&
-                                                     myContext.getPySettings().DICT_NEW_LINE_AFTER_LEFT_BRACE);
+      return !hasHangingIndent(myNode.getPsi())
+             && !(myNode.getElementType() == PyElementTypes.DICT_LITERAL_EXPRESSION &&
+                  myContext.getPySettings().DICT_NEW_LINE_AFTER_LEFT_BRACE)
+             && !(myNode.getElementType() == PyElementTypes.LIST_LITERAL_EXPRESSION &&
+                  myContext.getPySettings().LIST_NEW_LINE_AFTER_LEFT_BRACKET);
     }
     if (myNode.getElementType() == PyElementTypes.ARGUMENT_LIST) {
       if (!myContext.getSettings().ALIGN_MULTILINE_PARAMETERS_IN_CALLS || hasHangingIndent(myNode.getPsi())) {

python/python-syntax-core/src/com/jetbrains/python/formatter/PyCodeStyleSettings.java

@@ -9,6 +9,8 @@ import com.jetbrains.python.PySyntaxCoreBundle;
 import org.intellij.lang.annotations.MagicConstant;
 import org.jetbrains.annotations.NotNull;
 
+import static com.intellij.psi.codeStyle.CommonCodeStyleSettings.WRAP_AS_NEEDED;
+
 
 public class PyCodeStyleSettings extends CustomCodeStyleSettings {
 
@@ -71,6 +73,11 @@ public class PyCodeStyleSettings extends CustomCodeStyleSettings {
   public boolean DICT_NEW_LINE_AFTER_LEFT_BRACE = false;
   public boolean DICT_NEW_LINE_BEFORE_RIGHT_BRACE = false;
 
+  @CommonCodeStyleSettings.WrapConstant
+  public int LIST_WRAPPING = WRAP_AS_NEEDED;
+  public boolean LIST_NEW_LINE_AFTER_LEFT_BRACKET = false;
+  public boolean LIST_NEW_LINE_BEFORE_RIGHT_BRACKET = false;
+
   public int BLANK_LINES_AFTER_LOCAL_IMPORTS = 0;
   /**
    * Code style for most languages use continuation indent both for parameters in function definition and for arguments in function calls.

python/python-syntax-core/src/com/jetbrains/python/formatter/PythonFormattingModelBuilder.java

@@ -79,6 +79,11 @@ public class PythonFormattingModelBuilder implements FormattingModelBuilder, Cus
       .afterInside(LBRACE, DICT_LITERAL_EXPRESSION).spaceIf(pySettings.SPACE_WITHIN_BRACES, pySettings.DICT_NEW_LINE_AFTER_LEFT_BRACE)
       .beforeInside(RBRACE, DICT_LITERAL_EXPRESSION).spaceIf(pySettings.SPACE_WITHIN_BRACES, pySettings.DICT_NEW_LINE_BEFORE_RIGHT_BRACE)
 
+      .beforeInside(RBRACKET, LIST_LITERAL_EXPRESSION)
+      .spaceIf(commonSettings.SPACE_WITHIN_BRACKETS, pySettings.LIST_NEW_LINE_BEFORE_RIGHT_BRACKET)
+      .afterInside(LBRACKET, LIST_LITERAL_EXPRESSION)
+      .spaceIf(commonSettings.SPACE_WITHIN_BRACKETS, pySettings.LIST_NEW_LINE_AFTER_LEFT_BRACKET)
+
       .between(COMMA, RBRACE).spaceIf(pySettings.SPACE_WITHIN_BRACES || commonSettings.SPACE_AFTER_COMMA)
       .withinPair(LBRACE, RBRACE).spaceIf(pySettings.SPACE_WITHIN_BRACES)
 

python/python-syntax/resources/messages/PySyntaxBundle.properties

@@ -23,6 +23,7 @@ formatter.from.import.statements.force.parentheses.if.multiline=Force parenthese
 formatter.from.import.statements.force.comma.if.multline=Force trailing comma if multiline
 formatter.dictionary.literals=Dictionary literals
 formatter.hang.closing.brackets=Hang closing brackets
+formatter.list.literals=List literals
 
 formatter.imports.panel.optimize.imports=Optimize Imports
 formatter.imports.panel.sort.imports=Sort import statements

python/python-syntax/src/com/jetbrains/python/formatter/PyLanguageCodeStyleSettingsProvider.java

@@ -130,6 +130,14 @@ public final class PyLanguageCodeStyleSettingsProvider extends LanguageCodeStyle
                                 PySyntaxBundle.message("formatter.from.import.statements.force.comma.if.multline"),
                                 PySyntaxBundle.message("formatter.from.import.statements"));
 
+      consumer.showCustomOption(PyCodeStyleSettings.class, "LIST_WRAPPING",
+                                PySyntaxBundle.message("formatter.list.literals"), null, getInstance().WRAP_OPTIONS, WRAP_VALUES);
+      consumer.showCustomOption(PyCodeStyleSettings.class, "LIST_NEW_LINE_AFTER_LEFT_BRACKET",
+                                ApplicationBundle.message("wrapping.new.line.after.lbracket"), PySyntaxBundle.message("formatter.list.literals"));
+      consumer.showCustomOption(PyCodeStyleSettings.class, "LIST_NEW_LINE_BEFORE_RIGHT_BRACKET",
+                                ApplicationBundle.message("wrapping.rbracket.on.new.line"), PySyntaxBundle.message("formatter.list.literals"));
+
+
       consumer.showCustomOption(PyCodeStyleSettings.class, "DICT_WRAPPING",
                                 PySyntaxBundle.message("formatter.dictionary.literals"), null, getInstance().WRAP_OPTIONS, WRAP_VALUES);
       consumer.showCustomOption(PyCodeStyleSettings.class, "DICT_NEW_LINE_AFTER_LEFT_BRACE",

python/testData/formatter/listChopDownBracketsOnNewLine.py

@@ -0,0 +1,8 @@
+colors = ['red', 'green', 'blue', 'black', 'white', 'gray']
+
+
+def fooooooo():
+    colors = ['red', 'green', 'blue', 'black', 'white', 'gray']
+
+    def baaaaaar():
+        colors = ['red', 'green', 'blue', 'black', 'white', 'gray']

python/testData/formatter/listChopDownBracketsOnNewLine_after.py

@@ -0,0 +1,29 @@
+colors = [
+    'red',
+    'green',
+    'blue',
+    'black',
+    'white',
+    'gray'
+]
+
+
+def fooooooo():
+    colors = [
+        'red',
+        'green',
+        'blue',
+        'black',
+        'white',
+        'gray'
+    ]
+
+    def baaaaaar():
+        colors = [
+            'red',
+            'green',
+            'blue',
+            'black',
+            'white',
+            'gray'
+        ]

python/testData/formatter/listChopDownBracketsOnSameLine.py

@@ -0,0 +1,8 @@
+colors = ['red', 'green', 'blue', 'black', 'white', 'gray']
+
+
+def fooooooo():
+    colors = ['red', 'green', 'blue', 'black', 'white', 'gray']
+
+    def baaaaaar():
+        colors = ['red', 'green', 'blue', 'black', 'white', 'gray']

python/testData/formatter/listChopDownBracketsOnSameLine_after.py

@@ -0,0 +1,23 @@
+colors = ['red',
+          'green',
+          'blue',
+          'black',
+          'white',
+          'gray']
+
+
+def fooooooo():
+    colors = ['red',
+              'green',
+              'blue',
+              'black',
+              'white',
+              'gray']
+
+    def baaaaaar():
+        colors = ['red',
+                  'green',
+                  'blue',
+                  'black',
+                  'white',
+                  'gray']

python/testData/formatter/listWrappingBracketsOnNewLine.py

@@ -0,0 +1,8 @@
+colors = ['red', 'green', 'blue', 'black', 'white', 'gray']
+
+
+def fooooooo():
+    colors = ['red', 'green', 'blue', 'black', 'white', 'gray']
+
+    def baaaaaar():
+        colors = ['red', 'green', 'blue', 'black', 'white', 'gray']

python/testData/formatter/listWrappingBracketsOnNewLine_after.py

@@ -0,0 +1,17 @@
+colors = [
+    'red', 'green', 'blue', 'black',
+    'white', 'gray'
+]
+
+
+def fooooooo():
+    colors = [
+        'red', 'green', 'blue', 'black',
+        'white', 'gray'
+    ]
+
+    def baaaaaar():
+        colors = [
+            'red', 'green', 'blue',
+            'black', 'white', 'gray'
+        ]

python/testData/formatter/listWrappingBracketsOnSameLine.py

@@ -0,0 +1,8 @@
+colors = ['red', 'green', 'blue', 'black', 'white', 'gray']
+
+
+def fooooooo():
+    colors = ['red', 'green', 'blue', 'black', 'white', 'gray']
+
+    def baaaaaar():
+        colors = ['red', 'green', 'blue', 'black', 'white', 'gray']

python/testData/formatter/listWrappingBracketsOnSameLine_after.py

@@ -0,0 +1,12 @@
+colors = ['red', 'green', 'blue',
+          'black', 'white', 'gray']
+
+
+def fooooooo():
+    colors = ['red', 'green', 'blue',
+              'black', 'white', 'gray']
+
+    def baaaaaar():
+        colors = ['red', 'green',
+                  'blue', 'black',
+                  'white', 'gray']

python/testData/formatter/listWrappingDoNotAlign.py

@@ -0,0 +1,8 @@
+colors = ['red', 'green', 'blue', 'black', 'white', 'gray']
+
+
+def fooooooo():
+    colors = ['red', 'green', 'blue', 'black', 'white', 'gray']
+
+    def baaaaaar():
+        colors = ['red', 'green', 'blue', 'black', 'white', 'gray']

python/testData/formatter/listWrappingDoNotAlign_after.py

@@ -0,0 +1,26 @@
+colors = ['red',
+    'green',
+    'blue',
+    'black',
+    'white',
+    'gray'
+]
+
+
+def fooooooo():
+    colors = ['red',
+        'green',
+        'blue',
+        'black',
+        'white',
+        'gray'
+    ]
+
+    def baaaaaar():
+        colors = ['red',
+            'green',
+            'blue',
+            'black',
+            'white',
+            'gray'
+        ]

python/testData/formatter/listWrappingFirstBracketOnNewLine.py

@@ -0,0 +1,8 @@
+colors = ['red', 'green', 'blue', 'black', 'white', 'gray']
+
+
+def fooooooo():
+    colors = ['red', 'green', 'blue', 'black', 'white', 'gray']
+
+    def baaaaaar():
+        colors = ['red', 'green', 'blue', 'black', 'white', 'gray']

python/testData/formatter/listWrappingFirstBracketOnNewLine_after.py

@@ -0,0 +1,14 @@
+colors = [
+    'red', 'green', 'blue', 'black',
+    'white', 'gray']
+
+
+def fooooooo():
+    colors = [
+        'red', 'green', 'blue', 'black',
+        'white', 'gray']
+
+    def baaaaaar():
+        colors = [
+            'red', 'green', 'blue',
+            'black', 'white', 'gray']

python/testData/formatter/listWrappingLastBracketOnNewLine.py

@@ -0,0 +1,8 @@
+colors = ['red', 'green', 'blue', 'black', 'white', 'gray']
+
+
+def fooooooo():
+    colors = ['red', 'green', 'blue', 'black', 'white', 'gray']
+
+    def baaaaaar():
+        colors = ['red', 'green', 'blue', 'black', 'white', 'gray']

python/testData/formatter/listWrappingLastBracketOnNewLine_after.py

@@ -0,0 +1,15 @@
+colors = ['red', 'green', 'blue',
+          'black', 'white', 'gray'
+          ]
+
+
+def fooooooo():
+    colors = ['red', 'green', 'blue',
+              'black', 'white', 'gray'
+              ]
+
+    def baaaaaar():
+        colors = ['red', 'green',
+                  'blue', 'black',
+                  'white', 'gray'
+                  ]