code style settings to force new line after colon in single-clause and multi-clause compound statements

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

@@ -20,6 +20,7 @@ import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
 
@@ -124,11 +125,14 @@ public class PyBlock implements ASTBlock {
     Wrap wrap = null;
     Indent childIndent = Indent.getNoneIndent();
     Alignment childAlignment = null;
-    if (childType == PyElementTypes.STATEMENT_LIST || childType == PyElementTypes.IMPORT_ELEMENT) {
-      if (hasLineBreaksBefore(child, 1)) {
+    if (childType == PyElementTypes.STATEMENT_LIST) {
+      if (hasLineBreaksBefore(child, 1) || needLineBreakInStatement()) {
         childIndent = Indent.getNormalIndent();
       }
     }
+    else if (childType == PyElementTypes.IMPORT_ELEMENT && hasLineBreaksBefore(child, 1)) {
+      childIndent = Indent.getNormalIndent();
+    }
     if (ourListElementTypes.contains(parentType)) {
       // wrapping in non-parenthesized tuple expression is not allowed (PY-1792)
       if ((parentType != PyElementTypes.TUPLE_EXPRESSION || grandparentType == PyElementTypes.PARENTHESIZED_EXPRESSION) &&
@@ -383,16 +387,35 @@ public class PyBlock implements ASTBlock {
   @Nullable
   public Spacing getSpacing(Block child1, @NotNull Block child2) {
     if (child1 instanceof ASTBlock && child2 instanceof ASTBlock) {
-      final PsiElement psi1 = ((ASTBlock)child1).getNode().getPsi();
+      ASTNode node1 = ((ASTBlock)child1).getNode();
+      final PsiElement psi1 = node1.getPsi();
       final PsiElement psi2 = ((ASTBlock)child2).getNode().getPsi();
       if (psi1 instanceof PyImportStatementBase && psi2 instanceof PyImportStatementBase &&
           psi2.getCopyableUserData(IMPORT_GROUP_BEGIN) != null) {
         return Spacing.createSpacing(0, 0, 2, true, 1);
       }
+
+      if (node1.getElementType() == PyTokenTypes.COLON && psi2 instanceof PyStatementList) {
+        if (needLineBreakInStatement()) {
+          return Spacing.createSpacing(0, 0, 1, true, myContext.getSettings().KEEP_BLANK_LINES_IN_CODE);
+        }
+      }
     }
     return myContext.getSpacingBuilder().getSpacing(this, child1, child2);
   }
 
+  private boolean needLineBreakInStatement() {
+    PyStatement statement = PsiTreeUtil.getParentOfType(_node.getPsi(), PyStatement.class);
+    if (statement != null) {
+      Collection<PyStatementPart> parts = PsiTreeUtil.collectElementsOfType(statement, PyStatementPart.class);
+      if ((parts.size() == 1 && myContext.getPySettings().NEW_LINE_AFTER_COLON) ||
+          (parts.size() > 1 && myContext.getPySettings().NEW_LINE_AFTER_COLON_MULTI_CLAUSE)) {
+        return true;
+      }
+    }
+    return false;
+  }
+
   @NotNull
   public ChildAttributes getChildAttributes(int newChildIndex) {
     int statementListsBelow = 0;

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

@@ -20,6 +20,9 @@ public class PyCodeStyleSettings extends CustomCodeStyleSettings {
   public boolean ALIGN_COLLECTIONS_AND_COMPREHENSIONS = true;
   public boolean ALIGN_MULTILINE_IMPORTS = true;
 
+  public boolean NEW_LINE_AFTER_COLON = false;
+  public boolean NEW_LINE_AFTER_COLON_MULTI_CLAUSE = true;
+
   public PyCodeStyleSettings(CodeStyleSettings container) {
     super("Python", container);
   }

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

@@ -74,6 +74,10 @@ public class PyLanguageCodeStyleSettingsProvider extends LanguageCodeStyleSettin
                                    "WRAP_LONG_LINES",
                                    "ALIGN_MULTILINE_PARAMETERS",
                                    "ALIGN_MULTILINE_PARAMETERS_IN_CALLS");
+      consumer.showCustomOption(PyCodeStyleSettings.class, "NEW_LINE_AFTER_COLON", "Single-clause statements",
+                                "Force new line after colon");
+      consumer.showCustomOption(PyCodeStyleSettings.class, "NEW_LINE_AFTER_COLON_MULTI_CLAUSE", "Multi-clause statements",
+                                "Force new line after colon");
       consumer.showCustomOption(PyCodeStyleSettings.class, "ALIGN_COLLECTIONS_AND_COMPREHENSIONS", "Align when multiline",
                                 "Collections and Comprehensions");
       consumer.showCustomOption(PyCodeStyleSettings.class, "ALIGN_MULTILINE_IMPORTS", "Align when multiline",
@@ -126,7 +130,10 @@ public class PyLanguageCodeStyleSettingsProvider extends LanguageCodeStyleSettin
                                                 "xyzzy('long_string_constant1',\n" +
                                                 "    'long_string_constant2')\n" +
                                                 "attrs = [e.attr for e in\n" +
-                                                "    items]";
+                                                "    items]\n\n" +
+                                                "if True: pass\n\n" +
+                                                "try: pass\n" +
+                                                "finally: pass\n";
   @SuppressWarnings("FieldCanBeLocal")
   private static String INDENT_SETTINGS_PREVIEW = "def foo():\n" +
                                                   "    print 'bar'\n\n" +

python/testData/formatter/newLineAfterColon.py

@@ -0,0 +1 @@
+if True: pass

python/testData/formatter/newLineAfterColonMultiClause.py

@@ -0,0 +1,2 @@
+try: pass
+finally: pass

python/testData/formatter/newLineAfterColonMultiClause_after.py

@@ -0,0 +1,4 @@
+try:
+    pass
+finally:
+    pass

python/testData/formatter/newLineAfterColon_after.py

@@ -0,0 +1,2 @@
+if True:
+    pass

python/testData/refactoring/extractmethod/File.after.py

@@ -1,6 +1,8 @@
 def bar(base_new, self_new):
-    try: base_new.__init__(self_new)
-    except AttributeError: pass
+    try:
+        base_new.__init__(self_new)
+    except AttributeError:
+        pass
 
 
 def __init__(self):

python/testData/refactoring/extractmethod/MethodIndent.after.py

@@ -1,7 +1,9 @@
 class Foo(X, Y, Z):
     def bar(self, base_new):
-        try: base_new.__init__(self)
-        except AttributeError: pass
+        try:
+            base_new.__init__(self)
+        except AttributeError:
+            pass
 
     def __init__(self):
         for base in self__class__.__bases__:

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

@@ -281,6 +281,15 @@ public class PyFormatterTest extends PyTestCase {
     doTest();
   }
 
+  public void testNewLineAfterColon() {
+    settings().getCustomSettings(PyCodeStyleSettings.class).NEW_LINE_AFTER_COLON = true;
+    doTest();
+  }
+
+  public void testNewLineAfterColonMultiClause() {
+    doTest();
+  }
+
   private void doTest() {
     myFixture.configureByFile("formatter/" + getTestName(true) + ".py");
     ApplicationManager.getApplication().runWriteAction(new Runnable() {