PY-19100 Expand copied Python fragment to preserve original indentation

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

@@ -44,11 +44,50 @@ 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");
+  /**
+   * Keywords that start multiline block statements
+   */
+  private static final Set<String> START_KEYWORDS = ImmutableSet.of("async",
+                                                                    "def",
+                                                                    "class",
+                                                                    "with",
+                                                                    "if", "elif", "else",
+                                                                    "while", "for",
+                                                                    "try", "except", "finally");
 
   @Nullable
   @Override
   public String preprocessOnCopy(PsiFile file, int[] startOffsets, int[] endOffsets, String text) {
+    if (!CodeInsightSettings.getInstance().INDENT_TO_CARET_ON_PASTE || file.getLanguage() != PythonLanguage.getInstance()) {
+      return null;
+    }
+    // Expand copied text if it can cause indentation ambiguity
+    
+    // Text was selected with a single caret and might begin with a block statement 
+    if (startOffsets.length == 1 && endOffsets.length == 1 && fragmentBeginsWithBlockStatement(text)) {
+      final int start = startOffsets[0];
+      final int end = endOffsets[0];
+
+      final Document document = PsiDocumentManager.getInstance(file.getProject()).getDocument(file);
+      if (document != null) {
+        final int startLine = document.getLineNumber(start);
+        final int startLineOffset = getLineStartSafeOffset(document, startLine);
+        if (start != startLineOffset && startLine != document.getLineNumber(end)) {
+          final PsiElement keyword = file.findElementAt(start);
+          if (keyword != null && START_KEYWORDS.contains(keyword.getText())) {
+            final PyStatementListContainer block = PsiTreeUtil.getParentOfType(keyword, PyStatementListContainer.class);
+            // Statement body is in selection
+            if (block != null && end > block.getStatementList().getTextOffset()) {
+              final String linePrefix = document.getText(TextRange.create(startLineOffset, start));
+              if (StringUtil.isEmptyOrSpaces(linePrefix)) {
+                return linePrefix + text;
+              }
+            }
+          }
+        }
+      }
+    }
+    
     return null;
   }
 
@@ -78,7 +117,7 @@ public class PythonCopyPasteProcessor implements CopyPastePreProcessor {
     final PsiElement element = file.findElementAt(caretOffset);
     if (PsiTreeUtil.getParentOfType(element, PyStringLiteralExpression.class) != null) return text;
 
-    text = addLeadingSpacesToNormalizeSelection(project, file, text);
+    text = addLeadingSpacesToNormalizeSelection(project, text);
     final String indentText = getIndentText(file, document, caretOffset, lineNumber);
 
     final String line = document.getText(TextRange.create(lineStartOffset, lineEndOffset));
@@ -107,9 +146,8 @@ public class PythonCopyPasteProcessor implements CopyPastePreProcessor {
   }
 
   @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) {
+  private static String addLeadingSpacesToNormalizeSelection(@NotNull Project project, @NotNull String text) {
+    if (!fragmentBeginsWithBlockStatement(text)) {
       return text;
     }
 
@@ -134,6 +172,10 @@ public class PythonCopyPasteProcessor implements CopyPastePreProcessor {
     return text;
   }
 
+  private static boolean fragmentBeginsWithBlockStatement(@NotNull String text) {
+    return ContainerUtil.exists(START_KEYWORDS, keyword -> text.startsWith(keyword + " ") || text.startsWith(keyword + ":"));
+  }
+
   @NotNull
   private static String getIndentText(@NotNull final PsiFile file,
                                       @NotNull final Document document,

python/testData/copyPaste/AsyncFunctionWithBadSelection.after.py

@@ -0,0 +1,5 @@
+async def target():
+    x = 1
+
+
+y = 2

python/testData/copyPaste/AsyncFunctionWithBadSelection.dst.py

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

python/testData/copyPaste/AsyncFunctionWithBadSelection.src.py

@@ -0,0 +1,5 @@
+def f():
+    <selection>async def target():
+        x = 1
+
+    y = 2</selection>

python/testData/copyPaste/NonRectangleTopLevel.after.py

@@ -3,6 +3,5 @@ c = 1
 class A(object):
     def foo(self):
         pass
-
     def foo(self):
         pass

python/testData/copyPaste/TryBlockWithBadSelection.after.py

@@ -0,0 +1,3 @@
+try:
+    x = 1
+    y = 2

python/testData/copyPaste/TryBlockWithBadSelection.dst.py

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

python/testData/copyPaste/TryBlockWithBadSelection.src.py

@@ -0,0 +1,6 @@
+def f():
+    <selection>try:
+        x = 1
+        y = 2</selection>
+    finally:
+        pass

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

@@ -19,6 +19,7 @@ import com.intellij.codeInsight.CodeInsightSettings;
 import com.intellij.openapi.actionSystem.IdeActions;
 import com.intellij.psi.codeStyle.CommonCodeStyleSettings;
 import com.jetbrains.python.fixtures.PyTestCase;
+import com.jetbrains.python.psi.LanguageLevel;
 
 /**
  * @author yole
@@ -431,4 +432,14 @@ public class PyCopyPasteTest extends PyTestCase {
   public void testTopLevelIfStatementWithMultilineCondition() {
     doTest();
   }
+
+  // PY-19100
+  public void testTryBlockWithBadSelection() {
+    doTest();
+  }
+
+  // PY-19100
+  public void testAsyncFunctionWithBadSelection() {
+    runWithLanguageLevel(LanguageLevel.PYTHON35, this::doTest);
+  }
 }