@Override
  public PsiElement getTopLevelElement(PsiElement element) {
    PyPsiUtils.assertValid(element);
    final Ref<PsiFile> containingFileRef = Ref.create();
    ApplicationManager.getApplication().runReadAction(() -> containingFileRef.set(element.getContainingFile()));
    final PsiFile containingFile = containingFileRef.get();
    if (!(containingFile instanceof PyFile)) {
      return null;
    }
    List<PsiElement> parents = new ArrayList<>();
    PyDocStringOwner container = PsiTreeUtil.getParentOfType(element, PyDocStringOwner.class);
    while (container != null) {
      if (container instanceof PyFile) {
        break;
      }
      parents.add(0, container);
      container = PsiTreeUtil.getParentOfType(container, PyDocStringOwner.class);
    }
    for (PsiElement parent : parents) {
      if (parent instanceof PyFunction) {
        return parent;     // we don't display any nodes under functions
      }
    }
    if (parents.size() > 0) {
      return parents.get(parents.size() - 1);
    }
    return element.getContainingFile();
  }
@Override
  public PsiElement getTopLevelElement(PsiElement element)