@Override
  protected @Nullable PsiElement getStatementAtCaret(Editor editor, PsiFile psiFile) {
    PsiElement statementAtCaret = super.getStatementAtCaret(editor, psiFile);

    if (statementAtCaret instanceof PsiWhiteSpace) {
      return null;
    }
    if (statementAtCaret == null) {
      return null;
    }

    final PyStatementList statementList = PsiTreeUtil.getParentOfType(statementAtCaret, PyStatementList.class, false);
    if (statementList != null) {
      for (PyStatement statement : statementList.getStatements()) {
        if (PsiTreeUtil.isAncestor(statement, statementAtCaret, true)) {
          return statement;
        }
      }
    }
    else {
      final PyFile file = PsiTreeUtil.getParentOfType(statementAtCaret, PyFile.class, false);
      if (file != null) {
        for (PyStatement statement : file.getStatements()) {
          if (PsiTreeUtil.isAncestor(statement, statementAtCaret, true)) {
            return statement;
          }
        }
      }
    }
    return null;
  }
@Override
  protected @Nullable PsiElement getStatementAtCaret(Editor editor, PsiFile psiFile)