private static @Nullable PsiDirectory getPackageDir(@NotNull PsiElement elem) {
    if (elem instanceof PsiDirectory && PyUtil.isPackage(((PsiDirectory)elem), null)) {
      return (PsiDirectory)elem;
    }
    else if (elem instanceof PsiFile && PyUtil.isPackage(((PsiFile)elem))) {
      return ((PsiFile)elem).getParent();
    }
    return null;
  }
private static @Nullable PsiDirectory getPackageDir(@NotNull PsiElement elem)