imports are inserted after module-level dunder names to comply with PEP-8 (PY-23475)

GitOrigin-RevId: 9cad837e708f3c9e52abea59d9e239371cb515bc

python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/AddImportHelper.java

@@ -187,7 +187,7 @@ public class AddImportHelper {
     PsiElement feeler = insertParent.getFirstChild();
     if (feeler == null) return null;
     // skip initial comments and whitespace and try to get just below the last import stmt
-    boolean skippedOverImports = false;
+    boolean skippedOverStatements = false;
     boolean skippedOverDoc = false;
     PsiElement seeker = feeler;
     final boolean isInjected = InjectedLanguageManager.getInstance(feeler.getProject()).isInjectedFragment(feeler.getContainingFile());
@@ -206,7 +206,7 @@ public class AddImportHelper {
         }
         seeker = feeler;
         feeler = feeler.getNextSibling();
-        skippedOverImports = true;
+        skippedOverStatements = true;
       }
       else if (PsiTreeUtil.instanceOf(feeler, PsiWhiteSpace.class, PsiComment.class)) {
         seeker = feeler;
@@ -216,8 +216,13 @@ public class AddImportHelper {
         feeler = feeler.getNextSibling();
         seeker = feeler;
       }
+      else if (isAssignmentToModuleLevelDunderName(feeler)) {
+        feeler = feeler.getNextSibling();
+        seeker = feeler;
+        skippedOverStatements = true;
+      }
       // maybe we arrived at the doc comment stmt; skip over it, too
-      else if (!skippedOverImports && !skippedOverDoc && insertParent instanceof PyFile) {
+      else if (!skippedOverStatements && !skippedOverDoc && insertParent instanceof PyFile) {
         // this gives the literal; its parent is the expr seeker may have encountered
         final PsiElement docElem = DocStringUtil.findDocStringExpression((PyElement)insertParent);
         if (docElem != null && docElem.getParent() == feeler) {
@@ -572,6 +577,18 @@ public class AddImportHelper {
     }
   }
 
+  public static boolean isAssignmentToModuleLevelDunderName(@Nullable PsiElement element) {
+    if (element instanceof PyAssignmentStatement && PyUtil.isTopLevel(element)) {
+      PyAssignmentStatement statement = (PyAssignmentStatement)element;
+      PyExpression[] targets = statement.getTargets();
+      if (targets.length == 1) {
+        String name = targets[0].getName();
+        return name != null && PyUtil.isSpecialName(name);
+      }
+    }
+    return false;
+  }
+
   private static void addFileSystemItemImport(@NotNull PsiFileSystemItem target, @NotNull PsiFile file, @NotNull PyElement element) {
     final PsiFileSystemItem toImport = target.getParent();
     if (toImport == null) return;

python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyFileImpl.java

@@ -31,6 +31,7 @@ import com.jetbrains.python.PyNames;
 import com.jetbrains.python.PythonFileType;
 import com.jetbrains.python.PythonLanguage;
 import com.jetbrains.python.codeInsight.controlflow.ControlFlowCache;
+import com.jetbrains.python.codeInsight.imports.AddImportHelper;
 import com.jetbrains.python.documentation.docstrings.DocStringUtil;
 import com.jetbrains.python.psi.*;
 import com.jetbrains.python.psi.impl.references.PyReferenceImpl;
@@ -649,19 +650,37 @@ public class PyFileImpl extends PsiFileBase implements PyFile, PyExpression {
   public List<PyImportStatementBase> getImportBlock() {
     final List<PyImportStatementBase> result = new ArrayList<>();
     final PsiElement firstChild = getFirstChild();
-    final PyImportStatementBase firstImport;
+    PsiElement currentStatement;
     if (firstChild instanceof PyImportStatementBase) {
-      firstImport = (PyImportStatementBase)firstChild;
+      currentStatement = firstChild;
     }
     else {
-      firstImport = PsiTreeUtil.getNextSiblingOfType(firstChild, PyImportStatementBase.class);
+      currentStatement = PsiTreeUtil.getNextSiblingOfType(firstChild, PyImportStatementBase.class);
     }
-    if (firstImport != null) {
-      result.add(firstImport);
-      PsiElement nextImport = PyPsiUtils.getNextNonCommentSibling(firstImport, true);
-      while (nextImport instanceof PyImportStatementBase) {
-        result.add((PyImportStatementBase)nextImport);
-        nextImport = PyPsiUtils.getNextNonCommentSibling(nextImport, true);
+    if (currentStatement != null) {
+      // skip imports from future before module level dunders
+      final List<PyImportStatementBase> fromFuture = new ArrayList<>();
+      while (currentStatement instanceof PyFromImportStatement && ((PyFromImportStatement)currentStatement).isFromFuture()) {
+        fromFuture.add((PyImportStatementBase)currentStatement);
+        currentStatement = PyPsiUtils.getNextNonCommentSibling(currentStatement, true);
+      }
+
+      // skip module level dunders
+      boolean hasModuleLevelDunders = false;
+      while (AddImportHelper.isAssignmentToModuleLevelDunderName(currentStatement)) {
+        hasModuleLevelDunders = true;
+        currentStatement = PyPsiUtils.getNextNonCommentSibling(currentStatement, true);
+      }
+
+      // if there is an import from future and a module level-dunder between it and other imports,
+      // this import is not considered a part of the import block to avoid problems with "Optimize imports" and foldings
+      if (!hasModuleLevelDunders) {
+        result.addAll(fromFuture);
+      }
+      // collect imports
+      while (currentStatement instanceof PyImportStatementBase) {
+        result.add((PyImportStatementBase)currentStatement);
+        currentStatement = PyPsiUtils.getNextNonCommentSibling(currentStatement, true);
       }
     }
     return result;

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

@@ -15,6 +15,7 @@ import com.jetbrains.python.PyElementTypes;
 import com.jetbrains.python.PyTokenTypes;
 import com.jetbrains.python.PythonCodeStyleService;
 import com.jetbrains.python.PythonDialectsTokenSetProvider;
+import com.jetbrains.python.codeInsight.imports.AddImportHelper;
 import com.jetbrains.python.psi.*;
 import com.jetbrains.python.psi.impl.PyPsiUtils;
 import org.jetbrains.annotations.NotNull;
@@ -862,6 +863,11 @@ public class PyBlock implements ASTBlock {
         }
       }
 
+      if (psi2 instanceof PyImportStatementBase && AddImportHelper.isAssignmentToModuleLevelDunderName(psi1)) {
+        // blank line between module-level dunder name and import statement
+        return Spacing.createSpacing(0, 0, 2, settings.KEEP_LINE_BREAKS, settings.KEEP_BLANK_LINES_IN_DECLARATIONS);
+      }
+
       if ((PyElementTypes.CLASS_OR_FUNCTION.contains(childType1) && STATEMENT_OR_DECLARATION.contains(childType2)) ||
           STATEMENT_OR_DECLARATION.contains(childType1) && PyElementTypes.CLASS_OR_FUNCTION.contains(childType2)) {
         if (PyUtil.isTopLevel(psi1)) {

python/testData/addImport/moduleLevelDunder.after.py

@@ -0,0 +1,3 @@
+__author__ = "akniazev"
+
+from collections import OrderedDict

python/testData/addImport/moduleLevelDunder.py

@@ -0,0 +1 @@
+__author__ = "akniazev"

python/testData/addImport/moduleLevelDunderAndDocstring.after.py

@@ -0,0 +1,5 @@
+"""Top level docstring"""
+
+__author__ = "akniazev"
+
+from collections import OrderedDict

python/testData/addImport/moduleLevelDunderAndDocstring.py

@@ -0,0 +1,3 @@
+"""Top level docstring"""
+
+__author__ = "akniazev"

python/testData/addImport/moduleLevelDunderAndExistingImport.after.py

@@ -0,0 +1,4 @@
+__author__ = "akniazev"
+
+from collections import OrderedDict
+from sys import path

python/testData/addImport/moduleLevelDunderAndExistingImport.py

@@ -0,0 +1,3 @@
+__author__ = "akniazev"
+
+from sys import path

python/testData/addImport/moduleLevelDunderAndImportFromFuture.after.py

@@ -0,0 +1,5 @@
+from __future__ import absolute_import
+
+__author__ = "akniazev"
+
+from collections import OrderedDict

python/testData/addImport/moduleLevelDunderAndImportFromFuture.py

@@ -0,0 +1,3 @@
+from __future__ import absolute_import
+
+__author__ = "akniazev"

python/testData/formatter/moduleLevelDunderWithImports.py

@@ -0,0 +1,6 @@
+"""
+Docstring for file
+"""
+from __future__ import print_function
+__author__ = "akniazev"
+from collections import OrderedDict

python/testData/formatter/moduleLevelDunderWithImports_after.py

@@ -0,0 +1,8 @@
+"""
+Docstring for file
+"""
+from __future__ import print_function
+
+__author__ = "akniazev"
+
+from collections import OrderedDict

python/testData/optimizeImports/importFromFutureWithRegularImports.after.py

@@ -0,0 +1,12 @@
+from __future__ import print_function
+
+from collections import OrderedDict
+from datetime import date
+from datetime import time
+
+from foo import bar
+
+date(1, 1, 1)
+time(1)
+OrderedDict()
+bar()

python/testData/optimizeImports/importFromFutureWithRegularImports.py

@@ -0,0 +1,12 @@
+from datetime import date
+from sys import path
+from foo import bar
+from __future__ import print_function
+from collections import OrderedDict
+from datetime import time
+
+
+date(1, 1, 1)
+time(1)
+OrderedDict()
+bar()

python/testData/optimizeImports/moduleLevelDunder.after.py

@@ -0,0 +1,11 @@
+__author__ = "akniazev"
+
+from collections import OrderedDict
+from datetime import date, time
+
+from foo import bar
+
+date(1, 1, 1)
+time(1)
+OrderedDict()
+bar()

python/testData/optimizeImports/moduleLevelDunder.py

@@ -0,0 +1,13 @@
+__author__ = "akniazev"
+
+from datetime import date
+from sys import path
+from foo import bar
+from collections import OrderedDict
+from datetime import time
+
+
+date(1, 1, 1)
+time(1)
+OrderedDict()
+bar()

python/testData/optimizeImports/moduleLevelDunderWithImportFromFutureAbove.after.py

@@ -0,0 +1,13 @@
+from __future__ import print_function
+
+__author__ = "akniazev"
+
+from collections import OrderedDict
+from datetime import date, time
+
+from foo import bar
+
+date(1, 1, 1)
+time(1)
+OrderedDict()
+bar()

python/testData/optimizeImports/moduleLevelDunderWithImportFromFutureAbove.py

@@ -0,0 +1,15 @@
+from __future__ import print_function
+
+__author__ = "akniazev"
+
+from datetime import date
+from sys import path
+from foo import bar
+from collections import OrderedDict
+from datetime import time
+
+
+date(1, 1, 1)
+time(1)
+OrderedDict()
+bar()

python/testData/optimizeImports/moduleLevelDunderWithImportFromFutureBelow.after.py

@@ -0,0 +1,14 @@
+__author__ = "akniazev"
+
+from __future__ import print_function
+
+from collections import OrderedDict
+from datetime import date
+from datetime import time
+
+from foo import bar
+
+date(1, 1, 1)
+time(1)
+OrderedDict()
+bar()

python/testData/optimizeImports/moduleLevelDunderWithImportFromFutureBelow.py

@@ -0,0 +1,14 @@
+__author__ = "akniazev"
+
+from __future__ import print_function
+from datetime import date
+from sys import path
+from foo import bar
+from collections import OrderedDict
+from datetime import time
+
+
+date(1, 1, 1)
+time(1)
+OrderedDict()
+bar()

python/testSrc/com/jetbrains/python/PyAddImportTest.java

@@ -136,6 +136,26 @@ public class PyAddImportTest extends PyTestCase {
     );
   }
 
+  // PY-23475
+  public void testModuleLevelDunder() {
+    doAddFromImport("collections", "OrderedDict", BUILTIN);
+  }
+
+  // PY-23475
+  public void testModuleLevelDunderAndImportFromFuture() {
+    doAddFromImport("collections", "OrderedDict", BUILTIN);
+  }
+
+  // PY-23475
+  public void testModuleLevelDunderAndExistingImport(){
+    doAddFromImport("collections", "OrderedDict", BUILTIN);
+  }
+
+  // PY-23475
+  public void testModuleLevelDunderAndDocstring(){
+    doAddFromImport("collections", "OrderedDict", BUILTIN);
+  }
+
   private void doAddOrUpdateFromImport(final String path, final String name, final ImportPriority priority) {
     myFixture.configureByFile(getTestName(true) + ".py");
     WriteCommandAction.runWriteCommandAction(myFixture.getProject(), () -> {

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

@@ -964,4 +964,9 @@ public class PyFormatterTest extends PyTestCase {
   public void testSpacesAroundColonEqInAssignmentExpression() {
     runWithLanguageLevel(LanguageLevel.PYTHON38, this::doTest);
   }
+
+  // PY-23475
+  public void testModuleLevelDunderWithImports() {
+    doTest();
+  }
 }

python/testSrc/com/jetbrains/python/PyOptimizeImportsTest.java

@@ -354,6 +354,28 @@ public class PyOptimizeImportsTest extends PyTestCase {
     doTest();
   }
 
+  // PY-23475
+  public void testModuleLevelDunder() {
+    getPythonCodeStyleSettings().OPTIMIZE_IMPORTS_JOIN_FROM_IMPORTS_WITH_SAME_SOURCE = true;
+    doTest();
+  }
+
+  // PY-23475
+  public void testModuleLevelDunderWithImportFromFutureAbove() {
+    getPythonCodeStyleSettings().OPTIMIZE_IMPORTS_JOIN_FROM_IMPORTS_WITH_SAME_SOURCE = true;
+    doTest();
+  }
+
+  // PY-23475
+  public void testModuleLevelDunderWithImportFromFutureBelow() {
+    doTest();
+  }
+
+  // PY-23475
+  public void testImportFromFutureWithRegularImports() {
+    doTest();
+  }
+
   private void doMultiFileTest() {
     final String testName = getTestName(true);
     myFixture.copyDirectoryToProject(testName, "");