conflicts from cherrypick

GitOrigin-RevId: 0c5ac8858db88c73c09979309462129171489bab

python/python-ast/src/com/jetbrains/python/ast/PyAstFile.java

@@ -17,8 +17,11 @@ package com.jetbrains.python.ast;
 
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFile;
+import com.intellij.psi.util.PsiTreeUtil;
 import com.jetbrains.python.ast.controlFlow.AstScopeOwner;
 import com.jetbrains.python.ast.docstring.DocStringUtilCore;
+import com.jetbrains.python.ast.impl.PyPsiUtilsCore;
+import com.jetbrains.python.ast.impl.PyUtilCore;
 import com.jetbrains.python.psi.FutureFeature;
 import com.jetbrains.python.psi.LanguageLevel;
 import org.jetbrains.annotations.ApiStatus;
@@ -57,6 +60,48 @@ public interface PyAstFile extends PyAstElement, PsiFile, PyAstDocStringOwner, A
     return DocStringUtilCore.getDocStringValue(this);
   }
 
+  /**
+   * Returns the sequential list of import statements in the beginning of the file.
+   */
+  default List<? extends PyAstImportStatementBase> getImportBlock() {
+    final List<PyAstImportStatementBase> result = new ArrayList<>();
+    final PsiElement firstChild = getFirstChild();
+    PsiElement currentStatement;
+    if (firstChild instanceof PyAstImportStatementBase) {
+      currentStatement = firstChild;
+    }
+    else {
+      currentStatement = PsiTreeUtil.getNextSiblingOfType(firstChild, PyAstImportStatementBase.class);
+    }
+    if (currentStatement != null) {
+      // skip imports from future before module level dunders
+      final List<PyAstImportStatementBase> fromFuture = new ArrayList<>();
+      while (currentStatement instanceof PyAstFromImportStatement && ((PyAstFromImportStatement)currentStatement).isFromFuture()) {
+        fromFuture.add((PyAstImportStatementBase)currentStatement);
+        currentStatement = PyPsiUtilsCore.getNextNonCommentSibling(currentStatement, true);
+      }
+
+      // skip module level dunders
+      boolean hasModuleLevelDunders = false;
+      while (PyUtilCore.isAssignmentToModuleLevelDunderName(currentStatement)) {
+        hasModuleLevelDunders = true;
+        currentStatement = PyPsiUtilsCore.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 PyAstImportStatementBase) {
+        result.add((PyAstImportStatementBase)currentStatement);
+        currentStatement = PyPsiUtilsCore.getNextNonCommentSibling(currentStatement, true);
+      }
+    }
+    return result;
+  }
+
   @Override
   default @Nullable PyAstStringLiteralExpression getDocStringExpression() {
     return DocStringUtilCore.findDocStringExpression(this);

python/python-ast/src/com/jetbrains/python/ast/PyAstStringLiteralExpression.java

@@ -252,4 +252,12 @@ public interface PyAstStringLiteralExpression extends PyAstLiteralExpression, St
       }
     };
   }
+
+  /**
+   * @return true if this element has single string node and its type is {@link com.jetbrains.python.PyTokenTypes#DOCSTRING}
+   */
+  default boolean isDocString() {
+    final List<ASTNode> stringNodes = getStringNodes();
+    return stringNodes.size() == 1 && stringNodes.get(0).getElementType() == PyTokenTypes.DOCSTRING;
+  }
 }

python/python-psi-api/src/com/jetbrains/python/psi/PyFile.java

@@ -109,6 +109,7 @@ public interface PyFile extends PyAstFile, PyElement, PsiFile, PyDocStringOwner,
   /**
    * Returns the sequential list of import statements in the beginning of the file.
    */
+  @Override
   List<PyImportStatementBase> getImportBlock();
 
   @Override

python/python-psi-api/src/com/jetbrains/python/psi/PyStringLiteralExpression.java

@@ -48,6 +48,7 @@ public interface PyStringLiteralExpression extends PyAstStringLiteralExpression,
   /**
    * @return true if this element has single string node and its type is {@link com.jetbrains.python.PyTokenTypes#DOCSTRING}
    */
+  @Override
   boolean isDocString();
 
   /**

python/src/com/jetbrains/python/PythonFoldingBuilder.kt

@@ -36,7 +36,7 @@ class PythonFoldingBuilder : CustomFoldingBuilder(), DumbAware {
       return "import ..."
     }
     if (node.elementType === PyElementTypes.STRING_LITERAL_EXPRESSION) {
-      val stringLiteralExpression = node.psi as PyStringLiteralExpression
+      val stringLiteralExpression = node.psi as PyAstStringLiteralExpression
       val prefix = stringLiteralExpression.stringElements[0].prefix
       if (stringLiteralExpression.isDocString()) {
         val stringValue = stringLiteralExpression.stringValue.trim { it <= ' ' }
@@ -94,7 +94,7 @@ class PythonFoldingBuilder : CustomFoldingBuilder(), DumbAware {
   private fun appendDescriptors(node: ASTNode, descriptors: MutableList<FoldingDescriptor?>) {
     val elementType = node.elementType
     if (node.psi is PyAstFile) {
-      val imports = (node.psi as PyFile).importBlock
+      val imports = (node.psi as PyAstFile).importBlock
       if (imports.size > 1) {
         val firstImport: PyAstImportStatementBase = imports[0]
         val lastImport: PyAstImportStatementBase = imports[imports.size - 1]