static @Nullable PsiType getContextType(@NotNull PsiElement context) {
    return context instanceof PsiVariable var ? var.getType() :
           context instanceof PsiMethod method ? method.getReturnType() : null;
  }
static @Nullable PsiType getContextType(@NotNull PsiElement context)