PyType (@NotNull PyFunction function, @NotNull String parameterName) { final NumpyDocString docString = forFunction(function, function); if (docString != null) { String paramType = docString.getParamType(parameterName); // If parameter name starts with "p_", and we failed to obtain it from the docstring, // try to obtain parameter named without such prefix. if (paramType == null && parameterName.startsWith("p_")) { paramType = docString.getParamType(parameterName.substring(2)); } if (paramType != null) { if (isUfuncType(function, paramType)) { return getPsiFacade(function).parseTypeAnnotation("numbers.Number or numpy.core.multiarray.ndarray or collections.Iterable", function); } final PyType numpyDocType = parseNumpyDocType(function, paramType); if ("size".equals(parameterName)) { return getPsiFacade(function).createUnionType(Arrays.asList(numpyDocType, PyBuiltinCache.getInstance(function).getIntType())); } return numpyDocType; } } return null; }