PY-35936 Allow formatter to fix "PEP 8: E701 multiple statements on one line (colon)"

(cherry picked from commit e955950c3d09fee9d816147a9e951af6cee53ae0)

IJ-CR-16824

GitOrigin-RevId: 0fa973ad2868e9388515d36a4f57f4ae3a28c4e3

python/python-psi-impl/src/com/jetbrains/python/formatter/PyBlock.java

@@ -831,7 +831,7 @@ public class PyBlock implements ASTBlock {
             return getBlankLinesForOption(pySettings.BLANK_LINES_BEFORE_FIRST_METHOD);
           }
         }
-        if (childType1 == PyTokenTypes.COLON && needLineBreakInStatement()) {
+        if (childType1 == PyTokenTypes.COLON && (needLineBreakInStatement())) {
           return Spacing.createSpacing(0, 0, 1, true, settings.KEEP_BLANK_LINES_IN_CODE);
         }
       }
@@ -935,11 +935,16 @@ public class PyBlock implements ASTBlock {
   }
 
   private boolean needLineBreakInStatement() {
-    final PyStatement statement = PsiTreeUtil.getParentOfType(myNode.getPsi(), PyStatement.class);
-    if (statement != null) {
-      final Collection<PyStatementPart> parts = PsiTreeUtil.collectElementsOfType(statement, PyStatementPart.class);
-      return (parts.size() == 1 && myContext.getPySettings().NEW_LINE_AFTER_COLON) ||
-             (parts.size() > 1 && myContext.getPySettings().NEW_LINE_AFTER_COLON_MULTI_CLAUSE);
+    if (myNode.getPsi() instanceof PyStatementListContainer) {
+      final PyStatement statement = PsiTreeUtil.getParentOfType(myNode.getPsi(), PyStatement.class);
+      if (statement != null) {
+        final Collection<PyStatementPart> parts = PsiTreeUtil.collectElementsOfType(statement, PyStatementPart.class);
+        return (parts.size() == 1 && myContext.getPySettings().NEW_LINE_AFTER_COLON) ||
+               (parts.size() > 1 && myContext.getPySettings().NEW_LINE_AFTER_COLON_MULTI_CLAUSE);
+      }
+      else {
+        return myContext.getPySettings().NEW_LINE_AFTER_COLON;
+      }
     }
     return false;
   }

python/testData/formatter/pep8MultipleStatementsOnOneLine.py

@@ -0,0 +1,9 @@
+class Foo: pass
+
+if x > 5: y = 10
+
+def f(): pass
+
+def mul(x, y): return x * y
+
+with foo() as foo: pass

python/testData/formatter/pep8MultipleStatementsOnOneLine_after.py

@@ -0,0 +1,18 @@
+class Foo:
+    pass
+
+
+if x > 5:
+    y = 10
+
+
+def f():
+    pass
+
+
+def mul(x, y):
+    return x * y
+
+
+with foo() as foo:
+    pass

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

@@ -46,6 +46,12 @@ public class PyFormatterTest extends PyTestCase {
     doTest();
   }
 
+  // PY-35936
+  public void testPep8MultipleStatementsOnOneLine() {
+    getPythonCodeStyleSettings().NEW_LINE_AFTER_COLON = true;
+    doTest();
+  }
+
   public void testUnaryMinus() {
     doTest();
   }