@Nullable
public static PyType getType(@NotNull PsiElement resolved, @NotNull List<PyType> elementTypes) {
final String qualifiedName = getQualifiedName(resolved);
final List<Integer> paramListTypePositions = new ArrayList<>();
final List<Integer> ellipsisTypePositions = new ArrayList<>();
for (int i = 0; i < elementTypes.size(); i++) {
final PyType type = elementTypes.get(i);
if (type instanceof PyTypeParser.ParameterListType) {
paramListTypePositions.add(i);
} else if (type instanceof PyTypeParser.EllipsisType) {
ellipsisTypePositions.add(i);
}
}
if (!paramListTypePositions.isEmpty()) {
if (!("typing.Callable".equals(qualifiedName) && paramListTypePositions.equals(list(0)))) {
return null;
}
}
if (!ellipsisTypePositions.isEmpty()) {
if (!("typing.Callable".equals(qualifiedName) && ellipsisTypePositions.equals(list(0))
|| "typing.Tuple".equals(qualifiedName)
&& ellipsisTypePositions.equals(list(1))
&& elementTypes.size() == 2)) {
return null;
}
}
if ("typing.Union".equals(qualifiedName)) {
return PyUnionType.union(elementTypes);
}
if ("typing.Optional".equals(qualifiedName) && elementTypes.size() == 1) {
return PyUnionType.union(elementTypes.get(0), PyNoneType.INSTANCE);
}
if ("typing.Callable".equals(qualifiedName) && elementTypes.size() == 2) {
final PyTypeParser.ParameterListType paramList =
as(elementTypes.get(0), PyTypeParser.ParameterListType.class);
if (paramList != null) {
return new PyCallableTypeImpl(paramList.getCallableParameters(), elementTypes.get(1));
}
if (elementTypes.get(0) instanceof PyTypeParser.EllipsisType) {
return new PyCallableTypeImpl(null, elementTypes.get(1));
}
}
if ("typing.Tuple".equals(qualifiedName)) {
if (elementTypes.size() > 1 && elementTypes.get(1) instanceof PyTypeParser.EllipsisType) {
return PyTupleType.createHomogeneous(resolved, elementTypes.get(0));
}
return PyTupleType.create(resolved, elementTypes);
}
final PyType builtinCollection = getBuiltinCollection(resolved);
if (builtinCollection instanceof PyClassType) {
final PyClassType classType = (PyClassType) builtinCollection;
return new PyCollectionTypeImpl(classType.getPyClass(), false, elementTypes);
}
return null;
}
@Nullable
private static PyType getParameterizedType(
@NotNull PsiElement element, @NotNull Context context) {
if (element instanceof PySubscriptionExpression) {
final PySubscriptionExpression subscriptionExpr = (PySubscriptionExpression) element;
final PyExpression operand = subscriptionExpr.getOperand();
final PyExpression indexExpr = subscriptionExpr.getIndexExpression();
final PyType operandType = getType(operand, context);
if (operandType instanceof PyClassType) {
final PyClass cls = ((PyClassType) operandType).getPyClass();
final List<PyType> indexTypes = getIndexTypes(subscriptionExpr, context);
if (PyNames.TUPLE.equals(cls.getQualifiedName())) {
if (indexExpr instanceof PyTupleExpression) {
final PyExpression[] elements = ((PyTupleExpression) indexExpr).getElements();
if (elements.length == 2 && isEllipsis(elements[1])) {
return PyTupleType.createHomogeneous(element, indexTypes.get(0));
}
}
return PyTupleType.create(element, indexTypes);
} else if (indexExpr != null) {
return new PyCollectionTypeImpl(cls, false, indexTypes);
}
}
}
return null;
}
@NotNull
public static ParseResult parsePep484FunctionTypeComment(
@NotNull PsiElement anchor, @NotNull String text) {
final ForwardDeclaration<ParseResult, PyElementType> typeExpr = ForwardDeclaration.create();
final Function<Pair<ParseResult, List<ParseResult>>, ParseResult> toParamTypeList =
pair -> {
if (pair != null) {
final ParseResult first = pair.getFirst();
final List<ParseResult> second = pair.getSecond();
final List<PyType> itemTypes = new ArrayList<>();
ParseResult result = first;
itemTypes.add(result.getType());
for (ParseResult r : second) {
result = result.merge(r);
itemTypes.add(r.getType());
}
return result.withType(ParameterListType.fromParameterTypes(itemTypes));
}
return EMPTY_RESULT.withType(new ParameterListType(Collections.emptyList()));
};
final FunctionalParser<ParseResult, PyElementType> ellipsis =
op("...").map(token -> EMPTY_RESULT.withType(EllipsisType.INSTANCE)).named("ellipsis");
final FunctionalParser<ParseResult, PyElementType> classType =
token(IDENTIFIER)
.then(many(op(".").skipThen(token(IDENTIFIER))))
.map(new MakeSimpleType(anchor))
.cached()
.named("class-type");
final FunctionalParser<ParseResult, PyElementType> typeParamList =
op("[")
.skipThen(maybe(typeExpr.then(many(op(",").skipThen(typeExpr)))))
.thenSkip(op("]"))
.map(toParamTypeList)
.named("type-param-list");
final FunctionalParser<ParseResult, PyElementType> typeParam =
typeExpr.or(typeParamList).or(ellipsis).named("type-param");
final FunctionalParser<ParseResult, PyElementType> genericType =
classType
.thenSkip(op("["))
.then(typeParam)
.then(many(op(",").skipThen(typeParam)))
.thenSkip(op("]"))
.map(
value -> {
final Pair<ParseResult, ParseResult> firstPair = value.getFirst();
final ParseResult first = firstPair.getFirst();
final ParseResult second = firstPair.getSecond();
final List<ParseResult> third = value.getSecond();
final List<PyType> typesInBrackets = new ArrayList<>();
typesInBrackets.add(second.getType());
ParseResult result = first;
result = result.merge(second);
for (ParseResult r : third) {
typesInBrackets.add(r.getType());
result = result.merge(r);
}
final List<PyType> elementTypes =
third.isEmpty()
? Collections.singletonList(second.getType())
: typesInBrackets;
final PsiElement resolved = first.getElement();
if (resolved != null) {
final PyType typingType = PyTypingTypeProvider.getType(resolved, elementTypes);
if (typingType != null) {
return result.withType(typingType);
}
}
return EMPTY_RESULT;
})
.named("generic-type");
typeExpr.define(genericType.or(classType)).named("type-expr");
final FunctionalParser<ParseResult, PyElementType> paramType =
maybe(op("*"))
.then(maybe(op("*")))
.then(typeExpr)
.map(
pair -> {
final ParseResult paramResult = pair.getSecond();
final PyType type = paramResult.getType();
int starCount = 0;
if (pair.getFirst().getFirst() != null) {
starCount++;
}
if (pair.getFirst().getSecond() != null) {
starCount++;
}
if (starCount == 0) {
return paramResult;
} else if (starCount == 1) {
final PyClassType tupleType = PyTupleType.createHomogeneous(anchor, type);
if (tupleType != null) {
return paramResult.withType(tupleType);
}
return EMPTY_RESULT;
} else if (starCount == 2) {
final PyBuiltinCache builtinCache = PyBuiltinCache.getInstance(anchor);
final PyClassType dictType = builtinCache.getDictType();
if (dictType != null) {
final PyClass pyClass = dictType.getPyClass();
return paramResult.withType(
new PyCollectionTypeImpl(
pyClass, false, Arrays.asList(builtinCache.getStrType(), type)));
}
return EMPTY_RESULT;
}
return EMPTY_RESULT;
})
.named("param-type");
final FunctionalParser<ParseResult, PyElementType> paramTypes =
paramType.then(many(op(",").skipThen(paramType))).map(toParamTypeList).named("param-types");
final FunctionalParser<ParseResult, PyElementType> funcType =
op("(")
.skipThen(maybe(paramTypes.or(ellipsis)))
.thenSkip(op(")"))
.thenSkip(op("->"))
.then(typeExpr)
.map(
value -> {
final ParseResult paramsResult = value.getFirst();
final ParseResult returnResult = value.getSecond();
final List<PyCallableParameter> parameters;
ParseResult result = returnResult;
if (paramsResult != null) {
result = result.merge(paramsResult);
final ParameterListType paramTypesList =
as(paramsResult.getType(), ParameterListType.class);
if (paramTypesList != null) {
parameters = paramTypesList.getCallableParameters();
}
// ellipsis
else {
parameters = null;
}
} else {
parameters = Collections.emptyList();
result = returnResult;
}
return result.withType(
new PyCallableTypeImpl(parameters, returnResult.getType()));
})
.named("func-type");
final FunctionalParser<ParseResult, PyElementType> typeFile =
funcType.endOfInput().named("function-type-comment");
try {
return typeFile.parse(tokenize(text));
} catch (ParserException e) {
return EMPTY_RESULT;
}
}
