<commit_before>package com.jetbrains.python.codeInsight.editorActions.moveUpDown;

import com.intellij.codeInsight.CodeInsightUtilBase;
import com.intellij.codeInsight.editorActions.moveUpDown.LineMover;
import com.intellij.codeInsight.editorActions.moveUpDown.LineRange;
import com.intellij.openapi.diagnostic.Logger;
import com.intellij.openapi.editor.Document;
import com.intellij.openapi.editor.Editor;
import com.intellij.openapi.editor.LogicalPosition;
import com.intellij.openapi.editor.SelectionModel;
import com.intellij.openapi.util.Pair;
import com.intellij.psi.*;
import com.intellij.psi.impl.PsiDocumentManagerImpl;
import com.intellij.psi.util.PsiTreeUtil;
import com.jetbrains.python.psi.*;
import com.jetbrains.python.psi.impl.PyPsiUtils;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import java.util.ArrayList;
import java.util.Arrays;

/**
 * Created by IntelliJ IDEA.
 * Author: Alexey.Ivanov
 * Date:   23.04.2010
 * Time:   16:44:57
 */
public class StatementMover extends LineMover {
  private static final Logger LOG = Logger.getInstance("#com.jetbrains.python.codeInsight.editorActions.moveUpDown.StatementMover");

  private static PsiElement[] findStatementsInRange(PsiFile file, int startOffset, int endOffset) {
    PsiElement element1 = file.findElementAt(startOffset);
    PsiElement element2 = file.findElementAt(endOffset - 1);
    if (element1 instanceof PsiWhiteSpace) {
      startOffset = element1.getTextRange().getEndOffset();
      element1 = file.findElementAt(startOffset);
    }
    if (element2 instanceof PsiWhiteSpace) {
      endOffset = element2.getTextRange().getStartOffset();
      element2 = file.findElementAt(endOffset - 1);
    }
    if (element1 == null || element2 == null) {
      return PsiElement.EMPTY_ARRAY;
    }

    PsiElement parent = PsiTreeUtil.findCommonParent(element1, element2);
    if (parent == null) {
      return PsiElement.EMPTY_ARRAY;
    }

    while (true) {
      if (parent instanceof PyStatement) {
        parent = parent.getParent();
        break;
      }
      if (parent instanceof PyStatementList) {
        break;
      }
      if (parent == null || parent instanceof PsiFile) {
        return PsiElement.EMPTY_ARRAY;
      }
      parent = parent.getParent();
    }

    if (!parent.equals(element1)) {
      while (!parent.equals(element1.getParent())) {
        element1 = element1.getParent();
      }
    }

    if (!parent.equals(element2)) {
      while (!parent.equals(element2.getParent())) {
        element2 = element2.getParent();
      }
    }

    PsiElement[] children = parent.getChildren();
    ArrayList<PsiElement> array = new ArrayList<PsiElement>();
    boolean flag = false;
    for (PsiElement child : children) {
      if (child.equals(element1)) {
        flag = true;
      }
      if (flag && !(child instanceof PsiWhiteSpace)) {
        array.add(child);
      }
      if (child.equals(element2)) {
        break;
      }
    }

    for (PsiElement element : array) {
      if (!(element instanceof PyStatement || element instanceof PsiWhiteSpace || element instanceof PsiComment)) {
        return PsiElement.EMPTY_ARRAY;
      }
    }
    return array.toArray(new PsiElement[array.size()]);
  }

  private static boolean isNotValidStatementRange(Pair<PsiElement, PsiElement> range) {
    return range == null ||
           range.first == null ||
           range.second == null ||
           range.first.getParent() != range.second.getParent();
  }

  @Nullable
  private static LineRange expandLineRange(LineRange range, Editor editor, PsiFile file) {
    final SelectionModel selectionModel = editor.getSelectionModel();
    Pair<PsiElement, PsiElement> psiRange;
    if (selectionModel.hasSelection()) {
      final int startOffset = selectionModel.getSelectionStart();
      final int endOffset = selectionModel.getSelectionEnd();
      final PsiElement[] psiElements = findStatementsInRange(file, startOffset, endOffset);
      if (psiElements.length == 0) {
        return null;
      }
      psiRange = new Pair<PsiElement, PsiElement>(psiElements[0], psiElements[psiElements.length - 1]);
    }
    else {
      final int offset = editor.getCaretModel().getOffset();
      PsiElement element = file.findElementAt(offset);
      psiRange = new Pair<PsiElement, PsiElement>(element, element);
    }

    psiRange = new Pair<PsiElement, PsiElement>(PsiTreeUtil.getNonStrictParentOfType(psiRange.getFirst(), PyStatement.class),
                                                PsiTreeUtil.getNonStrictParentOfType(psiRange.getSecond(), PyStatement.class));
    if (psiRange.getFirst() == null || psiRange.getSecond() == null) {
      return null;
    }

    final PsiElement parent = PsiTreeUtil.findCommonParent(psiRange.getFirst(), psiRange.getSecond());
    final Pair<PsiElement, PsiElement> elementRange = getElementRange(parent, psiRange.getFirst(), psiRange.getSecond());
    if (isNotValidStatementRange(elementRange)) {
      return null;
    }

    if (elementRange.getFirst() == elementRange.getSecond() && elementRange.getFirst() instanceof PyPassStatement) {
      return null;
    }

    final int endOffset = elementRange.getSecond().getTextRange().getEndOffset();
    final Document document = editor.getDocument();
    if (endOffset > document.getTextLength()) {
      LOG.assertTrue(!PsiDocumentManager.getInstance(file.getProject()).isUncommited(document));
      LOG.assertTrue(PsiDocumentManagerImpl.checkConsistency(file, document));
    }

    int endLine;
    if (endOffset == document.getTextLength()) {
      endLine = document.getLineCount();
    }
    else {
      endLine = editor.offsetToLogicalPosition(endOffset).line + 1;
      endLine = Math.min(endLine, document.getLineCount());
    }
    endLine = Math.max(endLine, range.endLine);
    final int startLine = Math.min(range.startLine, editor.offsetToLogicalPosition(elementRange.getFirst().getTextOffset()).line);
    return new LineRange(startLine, endLine);
  }

  private @Nullable PyStatementList myStatementListToAddPass;
  private @Nullable PyStatementList myStatementListToRemovePass;
  private @NotNull PsiElement[] myElementsToDecreaseIndent;
  private @NotNull PsiElement[] myElementsToIncreaseIndent;

  @Override
  public boolean checkAvailable(@NotNull Editor editor,
                                @NotNull PsiFile file,
                                @NotNull MoveInfo info,
                                boolean down) {
    myStatementListToAddPass = null;
    myStatementListToRemovePass = null;
    myElementsToDecreaseIndent = PsiElement.EMPTY_ARRAY;
    myElementsToIncreaseIndent = PsiElement.EMPTY_ARRAY;
    if (!super.checkAvailable(editor, file, info, down)) {
      return false;
    }
    info.indentSource = true;
    final LineRange range = expandLineRange(info.toMove, editor, file);
    if (range == null) {
      return false;
    }

    info.toMove = range;
    final int startOffset = editor.logicalPositionToOffset(new LogicalPosition(range.startLine, 0));
    final int endOffset = editor.logicalPositionToOffset(new LogicalPosition(range.endLine, 0));
    final PsiElement[] statements = findStatementsInRange(file, startOffset, endOffset);
    final int length = statements.length;
    if (length == 0) {
      return false;
    }

    range.firstElement = statements[0];
    range.lastElement = statements[length - 1];
    final Document document = editor.getDocument();
    PyStatement statement;
    assert info.toMove2.startLine + 1 == info.toMove2.endLine;
    if (down) {
      statement = PsiTreeUtil.getNextSiblingOfType(range.lastElement, PyStatement.class);
    }
    else {
      statement = PsiTreeUtil.getPrevSiblingOfType(range.firstElement, PyStatement.class);
    }

    if (statement == null) {<commit_msg>Python statement mover works only in Python files<commit_after>package com.jetbrains.python.codeInsight.editorActions.moveUpDown;

import com.intellij.codeInsight.CodeInsightUtilBase;
import com.intellij.codeInsight.editorActions.moveUpDown.LineMover;
import com.intellij.codeInsight.editorActions.moveUpDown.LineRange;
import com.intellij.openapi.diagnostic.Logger;
import com.intellij.openapi.editor.Document;
import com.intellij.openapi.editor.Editor;
import com.intellij.openapi.editor.LogicalPosition;
import com.intellij.openapi.editor.SelectionModel;
import com.intellij.openapi.util.Pair;
import com.intellij.psi.*;
import com.intellij.psi.impl.PsiDocumentManagerImpl;
import com.intellij.psi.util.PsiTreeUtil;
import com.jetbrains.python.psi.*;
import com.jetbrains.python.psi.impl.PyPsiUtils;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import java.util.ArrayList;
import java.util.Arrays;

/**
 * Created by IntelliJ IDEA.
 * Author: Alexey.Ivanov
 * Date:   23.04.2010
 * Time:   16:44:57
 */
public class StatementMover extends LineMover {
  private static final Logger LOG = Logger.getInstance("#com.jetbrains.python.codeInsight.editorActions.moveUpDown.StatementMover");

  private static PsiElement[] findStatementsInRange(PsiFile file, int startOffset, int endOffset) {
    PsiElement element1 = file.findElementAt(startOffset);
    PsiElement element2 = file.findElementAt(endOffset - 1);
    if (element1 instanceof PsiWhiteSpace) {
      startOffset = element1.getTextRange().getEndOffset();
      element1 = file.findElementAt(startOffset);
    }
    if (element2 instanceof PsiWhiteSpace) {
      endOffset = element2.getTextRange().getStartOffset();
      element2 = file.findElementAt(endOffset - 1);
    }
    if (element1 == null || element2 == null) {
      return PsiElement.EMPTY_ARRAY;
    }

    PsiElement parent = PsiTreeUtil.findCommonParent(element1, element2);
    if (parent == null) {
      return PsiElement.EMPTY_ARRAY;
    }

    while (true) {
      if (parent instanceof PyStatement) {
        parent = parent.getParent();
        break;
      }
      if (parent instanceof PyStatementList) {
        break;
      }
      if (parent == null || parent instanceof PsiFile) {
        return PsiElement.EMPTY_ARRAY;
      }
      parent = parent.getParent();
    }

    if (!parent.equals(element1)) {
      while (!parent.equals(element1.getParent())) {
        element1 = element1.getParent();
      }
    }

    if (!parent.equals(element2)) {
      while (!parent.equals(element2.getParent())) {
        element2 = element2.getParent();
      }
    }

    PsiElement[] children = parent.getChildren();
    ArrayList<PsiElement> array = new ArrayList<PsiElement>();
    boolean flag = false;
    for (PsiElement child : children) {
      if (child.equals(element1)) {
        flag = true;
      }
      if (flag && !(child instanceof PsiWhiteSpace)) {
        array.add(child);
      }
      if (child.equals(element2)) {
        break;
      }
    }

    for (PsiElement element : array) {
      if (!(element instanceof PyStatement || element instanceof PsiWhiteSpace || element instanceof PsiComment)) {
        return PsiElement.EMPTY_ARRAY;
      }
    }
    return array.toArray(new PsiElement[array.size()]);
  }

  private static boolean isNotValidStatementRange(Pair<PsiElement, PsiElement> range) {
    return range == null ||
           range.first == null ||
           range.second == null ||
           range.first.getParent() != range.second.getParent();
  }

  @Nullable
  private static LineRange expandLineRange(LineRange range, Editor editor, PsiFile file) {
    final SelectionModel selectionModel = editor.getSelectionModel();
    Pair<PsiElement, PsiElement> psiRange;
    if (selectionModel.hasSelection()) {
      final int startOffset = selectionModel.getSelectionStart();
      final int endOffset = selectionModel.getSelectionEnd();
      final PsiElement[] psiElements = findStatementsInRange(file, startOffset, endOffset);
      if (psiElements.length == 0) {
        return null;
      }
      psiRange = new Pair<PsiElement, PsiElement>(psiElements[0], psiElements[psiElements.length - 1]);
    }
    else {
      final int offset = editor.getCaretModel().getOffset();
      PsiElement element = file.findElementAt(offset);
      psiRange = new Pair<PsiElement, PsiElement>(element, element);
    }

    psiRange = new Pair<PsiElement, PsiElement>(PsiTreeUtil.getNonStrictParentOfType(psiRange.getFirst(), PyStatement.class),
                                                PsiTreeUtil.getNonStrictParentOfType(psiRange.getSecond(), PyStatement.class));
    if (psiRange.getFirst() == null || psiRange.getSecond() == null) {
      return null;
    }

    final PsiElement parent = PsiTreeUtil.findCommonParent(psiRange.getFirst(), psiRange.getSecond());
    final Pair<PsiElement, PsiElement> elementRange = getElementRange(parent, psiRange.getFirst(), psiRange.getSecond());
    if (isNotValidStatementRange(elementRange)) {
      return null;
    }

    if (elementRange.getFirst() == elementRange.getSecond() && elementRange.getFirst() instanceof PyPassStatement) {
      return null;
    }

    final int endOffset = elementRange.getSecond().getTextRange().getEndOffset();
    final Document document = editor.getDocument();
    if (endOffset > document.getTextLength()) {
      LOG.assertTrue(!PsiDocumentManager.getInstance(file.getProject()).isUncommited(document));
      LOG.assertTrue(PsiDocumentManagerImpl.checkConsistency(file, document));
    }

    int endLine;
    if (endOffset == document.getTextLength()) {
      endLine = document.getLineCount();
    }
    else {
      endLine = editor.offsetToLogicalPosition(endOffset).line + 1;
      endLine = Math.min(endLine, document.getLineCount());
    }
    endLine = Math.max(endLine, range.endLine);
    final int startLine = Math.min(range.startLine, editor.offsetToLogicalPosition(elementRange.getFirst().getTextOffset()).line);
    return new LineRange(startLine, endLine);
  }

  private @Nullable PyStatementList myStatementListToAddPass;
  private @Nullable PyStatementList myStatementListToRemovePass;
  private @NotNull PsiElement[] myElementsToDecreaseIndent;
  private @NotNull PsiElement[] myElementsToIncreaseIndent;

  @Override
  public boolean checkAvailable(@NotNull Editor editor,
                                @NotNull PsiFile file,
                                @NotNull MoveInfo info,
                                boolean down) {
    myStatementListToAddPass = null;
    myStatementListToRemovePass = null;
    myElementsToDecreaseIndent = PsiElement.EMPTY_ARRAY;
    myElementsToIncreaseIndent = PsiElement.EMPTY_ARRAY;
    if (!(file instanceof PyFile)) {
      return false;
    }
    if (!super.checkAvailable(editor, file, info, down)) {
      return false;
    }
    info.indentSource = true;
    final LineRange range = expandLineRange(info.toMove, editor, file);
    if (range == null) {
      return false;
    }

    info.toMove = range;
    final int startOffset = editor.logicalPositionToOffset(new LogicalPosition(range.startLine, 0));
    final int endOffset = editor.logicalPositionToOffset(new LogicalPosition(range.endLine, 0));
    final PsiElement[] statements = findStatementsInRange(file, startOffset, endOffset);
    final int length = statements.length;
    if (length == 0) {
      return false;
    }

    range.firstElement = statements[0];
    range.lastElement = statements[length - 1];
    final Document document = editor.getDocument();
    PyStatement statement;
    assert info.toMove2.startLine + 1 == info.toMove2.endLine;
    if (down) {
      statement = PsiTreeUtil.getNextSiblingOfType(range.lastElement, PyStatement.class);
    }
    else {
      statement = PsiTreeUtil.getPrevSiblingOfType(range.firstElement, PyStatement.class);
    }

    if (statement == null) {
{'fraction_non_alphanumeric': 0.1073176934325132, 'fraction_numerical': 0.006357683260219976, 'mean_word_length': 3.476380193511668, 'pattern_counts': {'":': 0, '<': 19, '<?xml version=': 0, '>': 21, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3638094', 'n_tokens_mistral': 4594, 'n_tokens_neox': 4494, 'n_words': 921}