public static Map<PsiElement, Collection<PsiElement>> getCallees(@NotNull PyElement element) {
    final MultiMap<PsiElement, PsiElement> callees = MultiMap.createOrderedSet();

    final PyResolveContext resolveContext =
      PyResolveContext.implicitContext(TypeEvalContext.userInitiated(element.getProject(), element.getContainingFile()));

    final PyRecursiveElementVisitor visitor = new PyRecursiveElementVisitor() {
      @Override
      public void visitPyParameterList(@NotNull PyParameterList node) {
      }

      @Override
      public void visitPyLambdaExpression(@NotNull PyLambdaExpression node) {
      }

      @Override
      public void visitPyFunction(@NotNull PyFunction innerFunction) {
        for (PyParameter parameter : innerFunction.getParameterList().getParameters()) {
          PsiElement defaultValue = parameter.getDefaultValue();
          if (defaultValue != null) {
            defaultValue.accept(this);
          }
        }
      }

      @Override
      public void visitPyCallExpression(@NotNull PyCallExpression node) {
        super.visitPyCallExpression(node);

        StreamEx
          .of(node.multiResolveCalleeFunction(resolveContext))
          .select(PyFunction.class)
          .forEach(function -> callees.putValue(function, node));
      }
    };

    visitor.visitElement(element);

    return callees.freezeValues();
  }
public static Map<PsiElement, Collection<PsiElement>> getCallees(@NotNull PyElement element)