@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 getCallableType(@NotNull PsiElement resolved, @NotNull Context context) { if (resolved instanceof PySubscriptionExpression) { final PySubscriptionExpression subscriptionExpr = (PySubscriptionExpression) resolved; final PyExpression operand = subscriptionExpr.getOperand(); final Collection<String> operandNames = resolveToQualifiedNames(operand, context.getTypeContext()); if (operandNames.contains("typing.Callable")) { final PyExpression indexExpr = subscriptionExpr.getIndexExpression(); if (indexExpr instanceof PyTupleExpression) { final PyTupleExpression tupleExpr = (PyTupleExpression) indexExpr; final PyExpression[] elements = tupleExpr.getElements(); if (elements.length == 2) { final PyExpression parametersExpr = elements[0]; final PyExpression returnTypeExpr = elements[1]; if (parametersExpr instanceof PyListLiteralExpression) { final List<PyCallableParameter> parameters = new ArrayList<>(); final PyListLiteralExpression listExpr = (PyListLiteralExpression) parametersExpr; for (PyExpression argExpr : listExpr.getElements()) { parameters.add(new PyCallableParameterImpl(null, getType(argExpr, context))); } final PyType returnType = getType(returnTypeExpr, context); return new PyCallableTypeImpl(parameters, returnType); } if (isEllipsis(parametersExpr)) { return new PyCallableTypeImpl(null, getType(returnTypeExpr, context)); } } } } } return null; }
@NotNull private static List<PyType> getIndexTypes( @NotNull PySubscriptionExpression expression, @NotNull Context context) { final List<PyType> types = new ArrayList<>(); final PyExpression indexExpr = expression.getIndexExpression(); if (indexExpr instanceof PyTupleExpression) { final PyTupleExpression tupleExpr = (PyTupleExpression) indexExpr; for (PyExpression expr : tupleExpr.getElements()) { types.add(getType(expr, context)); } } else if (indexExpr != null) { types.add(getType(indexExpr, context)); } return types; }
@Nullable private static Ref<PyType> getOptionalType( @NotNull PsiElement element, @NotNull Context context) { if (element instanceof PySubscriptionExpression) { final PySubscriptionExpression subscriptionExpr = (PySubscriptionExpression) element; final PyExpression operand = subscriptionExpr.getOperand(); final Collection<String> operandNames = resolveToQualifiedNames(operand, context.getTypeContext()); if (operandNames.contains("typing.Optional")) { final PyExpression indexExpr = subscriptionExpr.getIndexExpression(); if (indexExpr != null) { final PyType type = getType(indexExpr, context); if (type != null) { return Ref.create(PyUnionType.union(type, PyNoneType.INSTANCE)); } } return Ref.create(); } } return null; }