@Nullable public Ref<PyType> getParameterType( @NotNull PyNamedParameter param, @NotNull PyFunction func, @NotNull TypeEvalContext context) { final PyAnnotation annotation = param.getAnnotation(); if (annotation != null) { // XXX: Requires switching from stub to AST final PyExpression value = annotation.getValue(); if (value != null) { final PyType type = getType(value, new Context(context)); if (type != null) { final PyType optionalType = getOptionalTypeFromDefaultNone(param, type, context); return Ref.create(optionalType != null ? optionalType : type); } } } final String paramComment = param.getTypeCommentAnnotation(); if (paramComment != null) { return Ref.create(getStringBasedType(paramComment, param, new Context(context))); } final String comment = func.getTypeCommentAnnotation(); if (comment != null) { final PyTypeParser.ParseResult result = PyTypeParser.parsePep484FunctionTypeComment(param, comment); final PyCallableType functionType = as(result.getType(), PyCallableType.class); if (functionType != null) { final List<PyCallableParameter> paramTypes = functionType.getParameters(context); // Function annotation of kind (...) -> Type if (paramTypes == null) { return Ref.create(); } final PyParameter[] funcParams = func.getParameterList().getParameters(); final int startOffset = omitFirstParamInTypeComment(func) ? 1 : 0; for (int paramIndex = 0; paramIndex < funcParams.length; paramIndex++) { if (funcParams[paramIndex] == param) { final int typeIndex = paramIndex - startOffset; if (typeIndex >= 0 && typeIndex < paramTypes.size()) { return Ref.create(paramTypes.get(typeIndex).getType(context)); } break; } } } } return null; }
@Nullable private static PyType getOptionalTypeFromDefaultNone( @NotNull PyNamedParameter param, @NotNull PyType type, @NotNull TypeEvalContext context) { final PyExpression defaultValue = param.getDefaultValue(); if (defaultValue != null) { final PyType defaultType = context.getType(defaultValue); if (defaultType instanceof PyNoneType) { return PyUnionType.union(type, defaultType); } } return null; }
private static Pair<String, String> getTypeAndDescription( @Nullable final String docString, @NotNull final PyNamedParameter followed) { StructuredDocString structuredDocString = DocStringUtil.parse(docString); String type = null; String desc = null; if (structuredDocString != null) { final String name = followed.getName(); type = structuredDocString.getParamType(name); desc = structuredDocString.getParamDescription(name); } return Pair.create(type, desc); }