public static @Nullable TextRange getExpressionRangeAtOffset(final Project project, final Document document, final int offset) {
    return ReadAction.compute(() -> {

      final PsiFile psiFile = PsiDocumentManager.getInstance(project).getPsiFile(document);
      if (psiFile != null) {
        PsiElement element = psiFile.findElementAt(offset);
        if (!(element instanceof PyExpression) || element instanceof PyLiteralExpression) {
          element = PsiTreeUtil.getParentOfType(element, PyExpression.class);
        }
        if (element instanceof PyLiteralExpression) {
          return null;
        }
        if (element != null && isSimpleEnough(element) && isExpression(project, document.getText(element.getTextRange()))) {
          return element.getTextRange();
        }
      }
      return null;
    });
  }
public static @Nullable TextRange getExpressionRangeAtOffset(final Project project, final Document document, final int offset)