Ref<PyType> (@NotNull PyFunction function, @NotNull PyCallSiteExpression callSite, @NotNull TypeEvalContext context) { if (isApplicable(function)) { final PyExpression callee = callSite instanceof PyCallExpression ? ((PyCallExpression)callSite).getCallee() : null; final NumpyDocString docString = forFunction(function, callee); if (docString != null) { final List<SectionField> returns = docString.getReturnFields(); final PyPsiFacade facade = getPsiFacade(function); return switch (returns.size()) { case 0 -> null; case 1 -> // Function returns single value Optional .ofNullable(returns.get(0).getType()) .filter(StringUtil::isNotEmpty) .map(typeName -> isUfuncType(function, typeName) ? facade.parseTypeAnnotation("T", function) : parseNumpyDocType(function, typeName)) .map(type -> PyTypingTypeProvider.toAsyncIfNeeded(function, type)) .map(Ref::create) .orElse(null); default -> { // Function returns a tuple final List<PyType> unionMembers = new ArrayList<>(); final List<PyType> members = new ArrayList<>(); for (int i = 0; i < returns.size(); i++) { final String memberTypeName = returns.get(i).getType(); final PyType returnType = StringUtil.isNotEmpty(memberTypeName) ? parseNumpyDocType(function, memberTypeName) : null; final boolean isOptional = StringUtil.isNotEmpty(memberTypeName) && memberTypeName.contains("optional"); if (isOptional && i != 0) { if (members.size() > 1) { unionMembers.add(facade.createTupleType(members, function)); } else if (members.size() == 1) { unionMembers.add(members.get(0)); } } members.add(returnType); if (i == returns.size() - 1 && isOptional) { unionMembers.add(facade.createTupleType(members, function)); } } final PyType type = unionMembers.isEmpty() ? facade.createTupleType(members, function) : facade.createUnionType(unionMembers); yield Ref.create(PyTypingTypeProvider.toAsyncIfNeeded(function, type)); } }; } } return null; }