@NotNull
private static List<PyGenericType> collectGenericTypes(
@NotNull PyClass cls, @NotNull Context context) {
boolean isGeneric = false;
for (PyClass ancestor : cls.getAncestorClasses(context.getTypeContext())) {
if (GENERIC_CLASSES.contains(ancestor.getQualifiedName())) {
isGeneric = true;
break;
}
}
if (isGeneric) {
final ArrayList<PyGenericType> results = new ArrayList<>();
// XXX: Requires switching from stub to AST
for (PyExpression expr : cls.getSuperClassExpressions()) {
if (expr instanceof PySubscriptionExpression) {
final PyExpression indexExpr = ((PySubscriptionExpression) expr).getIndexExpression();
if (indexExpr != null) {
for (PsiElement resolved : tryResolving(indexExpr, context.getTypeContext())) {
final PyGenericType genericType = getGenericType(resolved, context);
if (genericType != null) {
results.add(genericType);
}
}
}
}
}
return results;
}
return Collections.emptyList();
}
@Nullable
private static PyGenericType getGenericType(
@NotNull PsiElement element, @NotNull Context context) {
if (element instanceof PyCallExpression) {
final PyCallExpression assignedCall = (PyCallExpression) element;
final PyExpression callee = assignedCall.getCallee();
if (callee != null) {
final Collection<String> calleeQNames =
resolveToQualifiedNames(callee, context.getTypeContext());
if (calleeQNames.contains("typing.TypeVar")) {
final PyExpression[] arguments = assignedCall.getArguments();
if (arguments.length > 0) {
final PyExpression firstArgument = arguments[0];
if (firstArgument instanceof PyStringLiteralExpression) {
final String name = ((PyStringLiteralExpression) firstArgument).getStringValue();
if (name != null) {
return new PyGenericType(name, getGenericTypeBound(arguments, context));
}
}
}
}
}
}
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;
}
@Nullable
private static PyType getTypeForResolvedElement(
@NotNull PsiElement resolved, @NotNull Context context) {
if (context.getExpressionCache().contains(resolved)) {
// Recursive types are not yet supported
return null;
}
context.getExpressionCache().add(resolved);
try {
final PyType unionType = getUnionType(resolved, context);
if (unionType != null) {
return unionType;
}
final Ref<PyType> optionalType = getOptionalType(resolved, context);
if (optionalType != null) {
return optionalType.get();
}
final PyType callableType = getCallableType(resolved, context);
if (callableType != null) {
return callableType;
}
final PyType parameterizedType = getParameterizedType(resolved, context);
if (parameterizedType != null) {
return parameterizedType;
}
final PyType builtinCollection = getBuiltinCollection(resolved);
if (builtinCollection != null) {
return builtinCollection;
}
final PyType genericType = getGenericType(resolved, context);
if (genericType != null) {
return genericType;
}
final Ref<PyType> classType = getClassType(resolved, context.getTypeContext());
if (classType != null) {
return classType.get();
}
final PyType stringBasedType = getStringBasedType(resolved, context);
if (stringBasedType != null) {
return stringBasedType;
}
return null;
} finally {
context.getExpressionCache().remove(resolved);
}
}
@Nullable
private static PyType getType(@NotNull PyExpression expression, @NotNull Context context) {
final List<PyType> members = Lists.newArrayList();
for (PsiElement resolved : tryResolving(expression, context.getTypeContext())) {
members.add(getTypeForResolvedElement(resolved, context));
}
return PyUnionType.union(members);
}
@Nullable
private static PyType getUnionType(@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.Union")) {
return PyUnionType.union(getIndexTypes(subscriptionExpr, context));
}
}
return null;
}
@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;
}