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