@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;
}
@Nullable
private static PyType getStringBasedType(
@NotNull String contents, @NotNull PsiElement anchor, @NotNull Context context) {
final Project project = anchor.getProject();
final PyExpressionCodeFragmentImpl codeFragment =
new PyExpressionCodeFragmentImpl(project, "dummy.py", contents, false);
codeFragment.setContext(anchor.getContainingFile());
final PsiElement element = codeFragment.getFirstChild();
if (element instanceof PyExpressionStatement) {
final PyExpression expr = ((PyExpressionStatement) element).getExpression();
if (expr instanceof PyTupleExpression) {
final PyTupleExpression tupleExpr = (PyTupleExpression) expr;
final List<PyType> elementTypes =
ContainerUtil.map(
tupleExpr.getElements(), elementExpr -> getType(elementExpr, context));
return PyTupleType.create(anchor, elementTypes);
}
return getType(expr, context);
}
return null;
}