From 62476fa88fd50f5bb3c7d464f13ceb8e61cf048b Mon Sep 17 00:00:00 2001
From: Dmitry Trofimov <dmitry.trofimov@jetbrains.com>
Date: Tue, 4 Feb 2014 13:59:25 +0400
Subject: [PATCH] Fix alignment in list comprehension in dict (PY-10076).

---
 .../jetbrains/python/formatter/PyBlock.java   | 24 +++++++++++++++++++
 .../formatter/alignListComprehensionInDict.py |  5 ++++
 .../alignListComprehensionInDict_after.py     |  5 ++++
 .../com/jetbrains/python/PyFormatterTest.java |  4 ++++
 4 files changed, 38 insertions(+)
 create mode 100644 python/testData/formatter/alignListComprehensionInDict.py
 create mode 100644 python/testData/formatter/alignListComprehensionInDict_after.py

diff --git a/python/src/com/jetbrains/python/formatter/PyBlock.java b/python/src/com/jetbrains/python/formatter/PyBlock.java
index 1a867237eab3..ef18524dd3bc 100644
--- a/python/src/com/jetbrains/python/formatter/PyBlock.java
+++ b/python/src/com/jetbrains/python/formatter/PyBlock.java
@@ -141,6 +141,26 @@ public class PyBlock implements ASTBlock {
     Wrap wrap = null;
     Indent childIndent = Indent.getNoneIndent();
     Alignment childAlignment = null;
+
+    //First we are going to set-up possible alignment to parents
+    PyBlock p = !breaksAlignment(parentType) ? myParent : null;
+    while (p != null) {
+      ASTNode pNode = p.getNode();
+      if (ourListElementTypes.contains(pNode.getElementType())) {
+        if (needListAlignment(child) && !isEmptyList(_node.getPsi()) && !isInControlStatement()) {
+
+          childAlignment = p.getChildAlignment();
+          break;
+        }
+      }
+      if (!breaksAlignment(pNode.getElementType())) {
+        p = p.myParent;
+      }
+      else {
+        break;
+      }
+    }
+
     if (childType == PyElementTypes.STATEMENT_LIST) {
       if (hasLineBreaksBefore(child, 1) || needLineBreakInStatement()) {
         childIndent = Indent.getNormalIndent();
@@ -309,6 +329,10 @@ public class PyBlock implements ASTBlock {
     return new PyBlock(this, child, childAlignment, childIndent, wrap, myContext);
   }
 
+  private static boolean breaksAlignment(IElementType type) {
+    return type != PyElementTypes.BINARY_EXPRESSION;
+  }
+
   private static Alignment getAlignmentOfChild(PyBlock b, int childNum) {
     if (b.getSubBlocks().size() > childNum) {
       ChildAttributes attributes = b.getChildAttributes(childNum);
diff --git a/python/testData/formatter/alignListComprehensionInDict.py b/python/testData/formatter/alignListComprehensionInDict.py
new file mode 100644
index 000000000000..7b742792fc75
--- /dev/null
+++ b/python/testData/formatter/alignListComprehensionInDict.py
@@ -0,0 +1,5 @@
+def foo():
+        return {field.key: field for key, field in inspect.getmembers(instance)
+                if isinstance(field, QueryableAttribute)
+                   and isinstance(field.property, ColumnProperty)
+        or field.foreign_keys}
\ No newline at end of file
diff --git a/python/testData/formatter/alignListComprehensionInDict_after.py b/python/testData/formatter/alignListComprehensionInDict_after.py
new file mode 100644
index 000000000000..bcb4c2205c12
--- /dev/null
+++ b/python/testData/formatter/alignListComprehensionInDict_after.py
@@ -0,0 +1,5 @@
+def foo():
+    return {field.key: field for key, field in inspect.getmembers(instance)
+            if isinstance(field, QueryableAttribute)
+            and isinstance(field.property, ColumnProperty)
+            or field.foreign_keys}
\ No newline at end of file
diff --git a/python/testSrc/com/jetbrains/python/PyFormatterTest.java b/python/testSrc/com/jetbrains/python/PyFormatterTest.java
index f59fd49b2d03..39836892db77 100644
--- a/python/testSrc/com/jetbrains/python/PyFormatterTest.java
+++ b/python/testSrc/com/jetbrains/python/PyFormatterTest.java
@@ -394,6 +394,10 @@ public class PyFormatterTest extends PyTestCase {
     doTest();
   }
 
+  public void testAlignListComprehensionInDict() { //PY-10076
+    doTest();
+  }
+
   private void doTest() {
     doTest(false);
   }