PY-19100 Properly detect that the first line of pasted text needs to be additionally indented

It becomes necessary when user copies compound statement
staring selection at the beginning of the first keyword like
"def" instead of using structural selection. Because of that
we lose indentation of that first line and whole underlying block
structure becomes ambiguous, i.e. it's no longer clear which of the
following lines belong to the body of such function and which
were at the same or smaller block level initially. To handle it
we use simple heuristic that detects whether the body of copied
function is  indented too much and if it's the case we shift
the first line further to the right as if the etire first line
was copied in the first place.

Previously we compared indentation of the first two lines for this
purpose and it didn't work for various compound statements that may
have multiline header, e.g. function definition with individual
parameters on their own lines, "if" statement or "while" loop which
condition spans multiple lines etc. Now we first parse pasted
fragment and then take the indentation of the actual statement list.

python/src/com/jetbrains/python/editor/PythonCopyPasteProcessor.java

@@ -15,28 +15,27 @@
  */
 package com.jetbrains.python.editor;
 
+import com.google.common.collect.ImmutableSet;
 import com.intellij.codeInsight.CodeInsightSettings;
 import com.intellij.codeInsight.editorActions.CopyPastePreProcessor;
 import com.intellij.openapi.application.ApplicationManager;
 import com.intellij.openapi.editor.*;
 import com.intellij.openapi.project.Project;
 import com.intellij.openapi.util.TextRange;
-import com.intellij.openapi.util.text.CharFilter;
 import com.intellij.openapi.util.text.StringUtil;
 import com.intellij.psi.*;
-import com.intellij.psi.codeStyle.CodeStyleSettings;
-import com.intellij.psi.codeStyle.CodeStyleSettingsManager;
 import com.intellij.psi.tree.IElementType;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.util.ObjectUtils;
+import com.intellij.util.containers.ContainerUtil;
 import com.jetbrains.python.PyTokenTypes;
-import com.jetbrains.python.PythonFileType;
 import com.jetbrains.python.PythonLanguage;
 import com.jetbrains.python.psi.*;
+import com.jetbrains.python.psi.impl.PyExpressionCodeFragmentImpl;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
-import java.util.List;
+import java.util.Set;
 
 import static com.jetbrains.python.psi.PyUtil.as;
 
@@ -45,6 +44,8 @@ import static com.jetbrains.python.psi.PyUtil.as;
  */
 public class PythonCopyPasteProcessor implements CopyPastePreProcessor {
 
+  private static final Set<String> STATEMENT_WITH_HEADER_START_KEYWORDS = ImmutableSet.of("def", "class", "with", "if", "while", "for");
+
   @Nullable
   @Override
   public String preprocessOnCopy(PsiFile file, int[] startOffsets, int[] endOffsets, String text) {
@@ -61,15 +62,6 @@ public class PythonCopyPasteProcessor implements CopyPastePreProcessor {
     if (!CodeInsightSettings.getInstance().INDENT_TO_CARET_ON_PASTE || file.getLanguage() != PythonLanguage.getInstance()) {
       return text;
     }
-    final CodeStyleSettings codeStyleSettings = CodeStyleSettingsManager.getSettings(project);
-    final boolean useTabs = codeStyleSettings.useTabCharacter(PythonFileType.INSTANCE);
-    final int indentSize = codeStyleSettings.getIndentSize(PythonFileType.INSTANCE);
-    CharFilter NOT_INDENT_FILTER = new CharFilter() {
-      public boolean accept(char ch) {
-        return ch != (useTabs ? '\t' : ' ');
-      }
-    };
-    final String indentChar = useTabs ? "\t" : " ";
 
     final CaretModel caretModel = editor.getCaretModel();
     final SelectionModel selectionModel = editor.getSelectionModel();
@@ -86,7 +78,7 @@ public class PythonCopyPasteProcessor implements CopyPastePreProcessor {
     final PsiElement element = file.findElementAt(caretOffset);
     if (PsiTreeUtil.getParentOfType(element, PyStringLiteralExpression.class) != null) return text;
 
-    text = addLeadingSpaces(text, NOT_INDENT_FILTER, indentSize, indentChar);
+    text = addLeadingSpacesToNormalizeSelection(project, file, text);
     final String indentText = getIndentText(file, document, caretOffset, lineNumber);
 
     final String line = document.getText(TextRange.create(lineStartOffset, lineEndOffset));
@@ -107,25 +99,42 @@ public class PythonCopyPasteProcessor implements CopyPastePreProcessor {
       newText = text;
     }
 
+    final boolean useTabs = PyIndentUtil.areTabsUsedForIndentation(project);
     if (addLinebreak(text, line, useTabs) && selectionModel.getSelectionStart() == selectionModel.getSelectionEnd()) {
       newText += "\n";
     }
     return newText;
   }
 
-  private static String addLeadingSpaces(String text, final CharFilter filter, int indentSize, String indentChar) {
-    final List<String> strings = StringUtil.split(text, "\n", false);
-    if (strings.size() > 1) {
-      int firstLineIndent = StringUtil.findFirst(strings.get(0), filter);
-      int secondLineIndent = StringUtil.findFirst(strings.get(1), filter);
-      final int diff = secondLineIndent - firstLineIndent;
-      if (diff > indentSize) {
-        text = StringUtil.repeat(indentChar, diff - indentSize) + text;
-      }
+  @NotNull
+  private static String addLeadingSpacesToNormalizeSelection(@NotNull Project project, @NotNull PsiFile file, final @NotNull String text) {
+    boolean applicable = ContainerUtil.exists(STATEMENT_WITH_HEADER_START_KEYWORDS, keyword -> text.startsWith(keyword + " "));
+    if (!applicable) {
+      return text;
+    }
+
+    final PyExpressionCodeFragmentImpl fragment = new PyExpressionCodeFragmentImpl(project, "dummy.py", text, false);
+    //fragment.setContext(file);
+    final PyStatementListContainer statement = as(fragment.getFirstChild(), PyStatementListContainer.class);
+    if (statement == null) {
+      return text;
+    }
+
+    final String statementIndent = PyIndentUtil.getElementIndent(statement);
+    if (!statementIndent.isEmpty()) {
+      return text;
+    }
+    
+    final String indentStep = PyIndentUtil.getIndentFromSettings(project);
+    final String bodyIndent = PyIndentUtil.getElementIndent(statement.getStatementList());
+    final String expectedBodyIndent = statementIndent + indentStep;
+    if (bodyIndent.startsWith(expectedBodyIndent)) {
+      return bodyIndent.substring(0, bodyIndent.length() - indentStep.length()) + text;
     }
     return text;
   }
 
+  @NotNull
   private static String getIndentText(@NotNull final PsiFile file,
                                       @NotNull final Document document,
                                       int caretOffset,
@@ -194,7 +203,7 @@ public class PythonCopyPasteProcessor implements CopyPastePreProcessor {
   }
 
   private static boolean shouldPasteOnPreviousLine(@NotNull final PsiFile file, @NotNull String text, int caretOffset) {
-    final boolean useTabs = CodeStyleSettingsManager.getSettings(file.getProject()).useTabCharacter(PythonFileType.INSTANCE);
+    final boolean useTabs = PyIndentUtil.areTabsUsedForIndentation(file.getProject());
     final PsiElement nonWS = PyUtil.findNextAtOffset(file, caretOffset, PsiWhiteSpace.class);
     if (nonWS == null || text.endsWith("\n")) {
       return true;

python/src/com/jetbrains/python/psi/PyIndentUtil.java

@@ -131,7 +131,7 @@ public class PyIndentUtil {
     return statementList;
   } 
 
-  public static int getExpectedElementIndentSize(@NotNull PsiElement anchor) {
+  private static int getExpectedElementIndentSize(@NotNull PsiElement anchor) {
     int depth = 0;
     PyStatementList block = getAnchorStatementList(anchor);
     while (block != null) {
@@ -141,6 +141,15 @@ public class PyIndentUtil {
     return depth * getIndentSizeFromSettings(anchor.getProject());
   }
 
+  public static boolean areTabsUsedForIndentation(@NotNull Project project) {
+    final CodeStyleSettings codeStyleSettings = CodeStyleSettingsManager.getInstance(project).getCurrentSettings();
+    return codeStyleSettings.useTabCharacter(PythonFileType.INSTANCE);
+  }
+
+  public static char getIndentCharacter(@NotNull Project project) {
+    return areTabsUsedForIndentation(project) ? '\t' : ' ';
+  }
+
   /**
    * Returns indentation size configured in the Python code style settings.
    * 
@@ -153,14 +162,16 @@ public class PyIndentUtil {
   }
 
   /**
-   * Returns indentation configured in the Python code style settings as plain space character repeated number times specified there.
-   * Note that it doesn't take into account usage of tab characters that might be configured there as well.
+   * Returns indentation configured in the Python code style settings either as space character repeated number times specified there
+   * or a single tab character if tabs are set to use for indentation.
    * 
    * @see #getIndentSizeFromSettings(Project) 
+   * @see #areTabsUsedForIndentation(Project) 
    */
   @NotNull
   public static String getIndentFromSettings(@NotNull Project project) {
-    return StringUtil.repeatSymbol(' ', getIndentSizeFromSettings(project));
+    final boolean useTabs = areTabsUsedForIndentation(project);
+    return useTabs ? "\t" : StringUtil.repeatSymbol(' ', getIndentSizeFromSettings(project));
   }
 
   @NotNull

python/testData/copyPaste/TopLevelFunctionWithMultilineParameterList.after.py

@@ -0,0 +1,6 @@
+def hello(x,
+          y):
+    """
+    foo
+    """
+    return None

python/testData/copyPaste/TopLevelFunctionWithMultilineParameterList.dst.py

@@ -0,0 +1 @@
+<caret>

python/testData/copyPaste/TopLevelFunctionWithMultilineParameterList.src.py

@@ -0,0 +1,6 @@
+<selection>def hello(x,
+          y):
+    """
+    foo
+    """
+    return None</selection>

python/testData/copyPaste/TopLevelIfStatementWithMultilineCondition.after.py

@@ -0,0 +1,4 @@
+if (True or (True or
+                 False)):
+    x = 1
+    y = 2

python/testData/copyPaste/TopLevelIfStatementWithMultilineCondition.dst.py

@@ -0,0 +1 @@
+<caret>

python/testData/copyPaste/TopLevelIfStatementWithMultilineCondition.src.py

@@ -0,0 +1,4 @@
+<selection>if (True or (True or
+                 False)):
+    x = 1
+    y = 2</selection>

python/testSrc/com/jetbrains/python/PyCopyPasteTest.java

@@ -421,5 +421,14 @@ public class PyCopyPasteTest extends PyTestCase {
   public void testAmbiguousParentBlockMidIndent() {
     doTest();
   }
-  
+
+  // PY-19100
+  public void testTopLevelFunctionWithMultilineParameterList() {
+    doTest();
+  }
+
+  // PY-19100
+  public void testTopLevelIfStatementWithMultilineCondition() {
+    doTest();
+  }
 }