private @Nullable PsiElement findSmartStepIntoBaseExpression(@NotNull Document document, int line) {
    PsiFile file = PsiDocumentManager.getInstance(mySession.getProject()).getPsiFile(document);
    if (file == null) return null;

    PsiElement element = file.findElementAt(DocumentUtil.getFirstNonSpaceCharOffset(document, line));
    if (element == null) return null;

    // Allow multiline smart step into for argument list case.
    PsiElement argumentList = PsiTreeUtil.getParentOfType(element, PyArgumentList.class);
    if (argumentList != null) return argumentList.getParent();

    // Allow multiline smart step into for parenthesized multiline expressions.
    PsiElement parenthesizedExpression = PsiTreeUtil.getParentOfType(element, PyParenthesizedExpression.class);
    if (parenthesizedExpression != null) return parenthesizedExpression;

    // If it's not an argument list or parenthesized expression, simply find and return top expression at the line.
    PsiElement parent = element.getParent();
    int lineStartOffset = document.getLineStartOffset(line);
    int lineEndOffset = document.getLineEndOffset(line);
    while (parent != null &&  lineStartOffset <= parent.getTextOffset() && parent.getTextOffset() < lineEndOffset) {
      element = parent;
      parent = element.getParent();
    }

    if (element instanceof PyExpression) return element;

    Ref<PsiElement> psiElementRef = new Ref<>();

    element.accept(new PyRecursiveElementVisitor() {
      @Override
      public void visitPyBinaryExpression(@NotNull PyBinaryExpression node) {
        storeElement(node);
      }

      @Override
      public void visitPyCallExpression(@NotNull PyCallExpression node) {
        storeElement(node);
      }

      @Override
      public void visitPyComprehensionElement(@NotNull PyComprehensionElement node) {
        storeElement(node);
      }

      @Override
      public void visitPyElement(@NotNull PyElement node) {
        if (node instanceof PyDecorator) {
          storeElement(node);
        }
        super.visitPyElement(node);
      }

      @Override
      public void visitPyDecoratorList(@NotNull PyDecoratorList node) {
        storeElement(node);
        super.visitPyDecoratorList(node);
      }

      private void storeElement(PsiElement node) {
        if (psiElementRef.isNull()) {
          psiElementRef.set(node);
        }
      }
    });

    return psiElementRef.get();
  }
private @Nullable PsiElement findSmartStepIntoBaseExpression(@NotNull Document document, int line)