less broken formatting for literals in argument list (PY-6672)

python/src/com/jetbrains/python/formatter/PyBlock.java

@@ -160,7 +160,7 @@ public class PyBlock implements ASTBlock {
       }
     }
     else if (parentType == PyElementTypes.DICT_LITERAL_EXPRESSION || parentType == PyElementTypes.SET_LITERAL_EXPRESSION) {
-      if (childType == PyTokenTypes.RBRACE) {
+      if (childType == PyTokenTypes.RBRACE || !hasLineBreaksBefore(child, 1)) {
         childIndent = Indent.getNoneIndent();
       }
       else {
@@ -238,10 +238,14 @@ public class PyBlock implements ASTBlock {
       if (!mySettings.ALIGN_MULTILINE_PARAMETERS_IN_CALLS) {
         return false;
       }
+      if (child.getElementType() == PyTokenTypes.COMMA) {
+        return false;
+      }
       PyArgumentList argList = (PyArgumentList)_node.getPsi();
       if (argList != null) {
         PyExpression[] arguments = argList.getArguments();
-        return arguments.length > 1 || (arguments.length == 1 && PyPsiUtils.getNextComma(arguments[0].getNode()) != null);
+        return arguments.length > 1 || hasLineBreaksBefore(child, 1) || (
+          arguments.length == 1 && PyPsiUtils.getNextComma(arguments[0].getNode()) != null);
       }
       return false;
     }

python/testData/formatter/setLiteralInArgList.py

@@ -0,0 +1,3 @@
+self.assertEqual(
+    {"000000000000", "111111111111"},
+    foo['bar']['baz'])

python/testData/formatter/setLiteralInArgList_after.py

@@ -0,0 +1,3 @@
+self.assertEqual(
+    {"000000000000", "111111111111"},
+    foo['bar']['baz'])

python/testSrc/com/jetbrains/python/PyFormatterTest.java

@@ -4,6 +4,7 @@ import com.intellij.openapi.application.ApplicationManager;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFile;
 import com.intellij.psi.codeStyle.CodeStyleManager;
+import com.intellij.psi.codeStyle.CodeStyleSettings;
 import com.intellij.psi.codeStyle.CodeStyleSettingsManager;
 import com.jetbrains.python.fixtures.PyTestCase;
 import com.jetbrains.python.psi.LanguageLevel;
@@ -118,12 +119,12 @@ public class PyFormatterTest extends PyTestCase {
   }
 
   public void testSpaceInMethodDeclaration() {  // PY-4241
-    CodeStyleSettingsManager.getInstance().getSettings(myFixture.getProject()).SPACE_BEFORE_METHOD_PARENTHESES = true;
+    settings().SPACE_BEFORE_METHOD_PARENTHESES = true;
     doTest();
   }
 
   public void testOptionalAlignForMethodParameters() {  // PY-3995
-    CodeStyleSettingsManager.getInstance().getSettings(myFixture.getProject()).ALIGN_MULTILINE_PARAMETERS = false;
+    settings().ALIGN_MULTILINE_PARAMETERS = false;
     doTest();
   }
 
@@ -132,7 +133,7 @@ public class PyFormatterTest extends PyTestCase {
   }
 
   public void testAlignForMethodArguments() {  // PY-3995
-    CodeStyleSettingsManager.getInstance().getSettings(myFixture.getProject()).ALIGN_MULTILINE_PARAMETERS_IN_CALLS = true;
+    settings().ALIGN_MULTILINE_PARAMETERS_IN_CALLS = true;
     doTest();
   }
 
@@ -152,6 +153,11 @@ public class PyFormatterTest extends PyTestCase {
     doTest();
   }
 
+  public void testSetLiteralInArgList() {  // PY-6672
+    settings().ALIGN_MULTILINE_PARAMETERS_IN_CALLS = true;
+    doTest();
+  }
+
   public void testPsiFormatting() { // IDEA-69724
     String initial =
       "def method_name(\n" +
@@ -190,4 +196,8 @@ public class PyFormatterTest extends PyTestCase {
     });
     myFixture.checkResultByFile("formatter/" + getTestName(true) + "_after.py");
   }
+
+  private CodeStyleSettings settings() {
+    return CodeStyleSettingsManager.getInstance().getSettings(myFixture.getProject());
+  }
 }