@Override
public PyType getReferenceType(
@NotNull PsiElement referenceTarget, TypeEvalContext context, @Nullable PsiElement anchor) {
if (referenceTarget instanceof PyTargetExpression) {
final PyTargetExpression target = (PyTargetExpression) referenceTarget;
if (context.maySwitchToAST(target)) {
// XXX: Requires switching from stub to AST
final PyAnnotation annotation = target.getAnnotation();
if (annotation != null) {
final PyExpression value = annotation.getValue();
if (value != null) {
return getType(value, new Context(context));
}
return null;
}
}
final String comment = target.getTypeCommentAnnotation();
if (comment != null) {
final PyType type = getStringBasedType(comment, referenceTarget, new Context(context));
if (type instanceof PyTupleType) {
final PyTupleExpression tupleExpr =
PsiTreeUtil.getParentOfType(target, PyTupleExpression.class);
if (tupleExpr != null) {
return PyTypeChecker.getTargetTypeFromTupleAssignment(
target, tupleExpr, (PyTupleType) type);
}
}
return type;
}
}
return null;
}
@Nullable
@Override
public PyType getReturnType(@NotNull TypeEvalContext context, @NotNull TypeEvalContext.Key key) {
for (PyTypeProvider typeProvider : Extensions.getExtensions(PyTypeProvider.EP_NAME)) {
final PyType returnType = typeProvider.getReturnType(this, context);
if (returnType != null) {
returnType.assertValid(typeProvider.toString());
return returnType;
}
}
if (context.maySwitchToAST(this)
&& LanguageLevel.forElement(this).isAtLeast(LanguageLevel.PYTHON30)) {
PyAnnotation anno = getAnnotation();
if (anno != null) {
PyClass pyClass = anno.resolveToClass();
if (pyClass != null) {
return new PyClassTypeImpl(pyClass, false);
}
}
}
final PyType docStringType = getReturnTypeFromDocString();
if (docStringType != null) {
docStringType.assertValid("from docstring");
return docStringType;
}
if (context.allowReturnTypes(this)) {
final Ref<? extends PyType> yieldTypeRef = getYieldStatementType(context);
if (yieldTypeRef != null) {
return yieldTypeRef.get();
}
return getReturnStatementType(context);
}
return null;
}
@Nullable
@Override
public Ref<PyType> getReturnType(@NotNull PyCallable callable, @NotNull TypeEvalContext context) {
if (callable instanceof PyFunction) {
final PyFunction function = (PyFunction) callable;
final PyAnnotation annotation = function.getAnnotation();
if (annotation != null) {
// XXX: Requires switching from stub to AST
final PyExpression value = annotation.getValue();
if (value != null) {
final PyType type = getType(value, new Context(context));
return type != null ? Ref.create(type) : null;
}
}
final PyType constructorType = getGenericConstructorType(function, new Context(context));
if (constructorType != null) {
return Ref.create(constructorType);
}
final String comment = function.getTypeCommentAnnotation();
if (comment != null) {
final PyTypeParser.ParseResult result =
PyTypeParser.parsePep484FunctionTypeComment(callable, comment);
final PyCallableType funcType = as(result.getType(), PyCallableType.class);
if (funcType != null) {
return Ref.create(funcType.getReturnType(context));
}
}
}
return null;
}
@Nullable
private PyType getGenericReturnType(
@NotNull TypeEvalContext typeEvalContext, @Nullable PyQualifiedExpression callSite) {
if (typeEvalContext.maySwitchToAST(this)
&& LanguageLevel.forElement(this).isAtLeast(LanguageLevel.PYTHON30)) {
PyAnnotation anno = getAnnotation();
if (anno != null) {
PyClass pyClass = anno.resolveToClass();
if (pyClass != null) {
return new PyClassTypeImpl(pyClass, false);
}
}
}
for (PyTypeProvider typeProvider : Extensions.getExtensions(PyTypeProvider.EP_NAME)) {
final PyType returnType = typeProvider.getReturnType(this, callSite, typeEvalContext);
if (returnType != null) {
returnType.assertValid(typeProvider.toString());
return returnType;
}
}
final PyType docStringType = getReturnTypeFromDocString();
if (docStringType != null) {
docStringType.assertValid("from docstring");
return docStringType;
}
if (typeEvalContext.allowReturnTypes(this)) {
final PyType yieldType = getYieldStatementType(typeEvalContext);
if (yieldType != null) {
return yieldType;
}
return getReturnStatementType(typeEvalContext);
}
return null;
}
@Nullable
public Ref<PyType> getParameterType(
@NotNull PyNamedParameter param, @NotNull PyFunction func, @NotNull TypeEvalContext context) {
final PyAnnotation annotation = param.getAnnotation();
if (annotation != null) {
// XXX: Requires switching from stub to AST
final PyExpression value = annotation.getValue();
if (value != null) {
final PyType type = getType(value, new Context(context));
if (type != null) {
final PyType optionalType = getOptionalTypeFromDefaultNone(param, type, context);
return Ref.create(optionalType != null ? optionalType : type);
}
}
}
final String paramComment = param.getTypeCommentAnnotation();
if (paramComment != null) {
return Ref.create(getStringBasedType(paramComment, param, new Context(context)));
}
final String comment = func.getTypeCommentAnnotation();
if (comment != null) {
final PyTypeParser.ParseResult result =
PyTypeParser.parsePep484FunctionTypeComment(param, comment);
final PyCallableType functionType = as(result.getType(), PyCallableType.class);
if (functionType != null) {
final List<PyCallableParameter> paramTypes = functionType.getParameters(context);
// Function annotation of kind (...) -> Type
if (paramTypes == null) {
return Ref.create();
}
final PyParameter[] funcParams = func.getParameterList().getParameters();
final int startOffset = omitFirstParamInTypeComment(func) ? 1 : 0;
for (int paramIndex = 0; paramIndex < funcParams.length; paramIndex++) {
if (funcParams[paramIndex] == param) {
final int typeIndex = paramIndex - startOffset;
if (typeIndex >= 0 && typeIndex < paramTypes.size()) {
return Ref.create(paramTypes.get(typeIndex).getType(context));
}
break;
}
}
}
}
return null;
}