@NotNull
  private static PsiElement getMultiCalleeElementToHighlight(@NotNull PyCallSiteExpression callSite) {
    if (callSite instanceof PyCallExpression call) {
      final PyArgumentList argumentList = call.getArgumentList();

      final PsiElement result = Optional
        .ofNullable(argumentList)
        .map(PyArgumentList::getArguments)
        .filter(arguments -> arguments.length == 1)
        .<PsiElement>map(arguments -> arguments[0])
        .orElse(argumentList);

      return ObjectUtils.notNull(result, callSite);
    }
    else if (callSite instanceof PySubscriptionExpression) {
      return ObjectUtils.notNull(((PySubscriptionExpression)callSite).getIndexExpression(), callSite);
    }
    else {
      return callSite;
    }
  }
@NotNull
  private static PsiElement getMultiCalleeElementToHighlight(@NotNull PyCallSiteExpression callSite)