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