@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);
}
