Examples with NumpyDocString com.jetbrains.python.documentation.docstrings.NumpyDocString used on opensource projects

Example 1 with NumpyDocString

use of com.jetbrains.python.documentation.docstrings.NumpyDocString in project intellij-community by JetBrains.

the class NumpyDocStringTypeProvider method getCallType.

@Nullable
@Override
public Ref<PyType> getCallType(@NotNull PyFunction function, @Nullable 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);
            switch(returns.size()) {
                case 0:
                    return null;
                case 1:
                    // Function returns single value
                    return Optional.ofNullable(returns.get(0).getType()).filter(StringUtil::isNotEmpty).map(typeName -> isUfuncType(function, typeName) ? facade.parseTypeAnnotation("T", function) : parseNumpyDocType(function, typeName)).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);
                    return Ref.create(type);
            }
        }
    }
    return null;
}

Example 2 with NumpyDocString

use of com.jetbrains.python.documentation.docstrings.NumpyDocString in project intellij-community by JetBrains.

the class NumpyDocStringTypeProvider method forFunction.

@Nullable
private static NumpyDocString forFunction(@NotNull PyFunction function, @Nullable PsiElement reference, @Nullable String knownSignature) {
    String docString = function.getDocStringValue();
    if (docString == null && PyNames.INIT.equals(function.getName())) {
        // Docstring for constructor can be found in the docstring of class
        PyClass cls = function.getContainingClass();
        if (cls != null) {
            docString = cls.getDocStringValue();
        }
    }
    if (docString != null) {
        final NumpyDocString parsed = (NumpyDocString) DocStringUtil.parseDocString(DocStringFormat.NUMPY, docString);
        if (parsed.getReturnFields().isEmpty() && parsed.getParameterFields().isEmpty()) {
            return null;
        }
        String signature = parsed.getSignature();
        String redirect = findRedirect(parsed.getLines());
        if (redirect != null && reference != null) {
            PyFunction resolvedFunction = resolveRedirectToFunction(redirect, reference);
            if (resolvedFunction != null) {
                return forFunction(resolvedFunction, reference, knownSignature != null ? knownSignature : signature);
            }
        }
        return parsed;
    }
    return null;
}