@Override
public void visitPyReturnStatement(PyReturnStatement node) {
if (PsiTreeUtil.getParentOfType(node, ScopeOwner.class, true) == myFunction) {
final PyExpression expr = node.getExpression();
PyType returnType;
returnType = expr == null ? PyNoneType.INSTANCE : myContext.getType(expr);
if (!myHasReturns) {
myResult = returnType;
myHasReturns = true;
} else {
myResult = PyUnionType.union(myResult, returnType);
}
}
}
@Nullable
private Ref<? extends PyType> getYieldStatementType(@NotNull final TypeEvalContext context) {
Ref<PyType> elementType = null;
final PyBuiltinCache cache = PyBuiltinCache.getInstance(this);
final PyStatementList statements = getStatementList();
final Set<PyType> types = new LinkedHashSet<>();
statements.accept(
new PyRecursiveElementVisitor() {
@Override
public void visitPyYieldExpression(PyYieldExpression node) {
final PyExpression expr = node.getExpression();
final PyType type = expr != null ? context.getType(expr) : null;
if (node.isDelegating() && type instanceof PyCollectionType) {
final PyCollectionType collectionType = (PyCollectionType) type;
// TODO: Select the parameter types that matches T in Iterable[T]
final List<PyType> elementTypes = collectionType.getElementTypes(context);
types.add(elementTypes.isEmpty() ? null : elementTypes.get(0));
} else {
types.add(type);
}
}
@Override
public void visitPyFunction(PyFunction node) {
// Ignore nested functions
}
});
final int n = types.size();
if (n == 1) {
elementType = Ref.create(types.iterator().next());
} else if (n > 0) {
elementType = Ref.create(PyUnionType.union(types));
}
if (elementType != null) {
final PyClass generator = cache.getClass(PyNames.FAKE_GENERATOR);
if (generator != null) {
final List<PyType> parameters =
Arrays.asList(elementType.get(), null, getReturnStatementType(context));
return Ref.create(new PyCollectionTypeImpl(generator, false, parameters));
}
}
if (!types.isEmpty()) {
return Ref.create(null);
}
return null;
}
@Nullable
private Ref<? extends PyType> getYieldStatementType(@NotNull final TypeEvalContext context) {
Ref<PyType> elementType = null;
final PyBuiltinCache cache = PyBuiltinCache.getInstance(this);
final PyStatementList statements = getStatementList();
final Set<PyType> types = new LinkedHashSet<PyType>();
if (statements != null) {
statements.accept(
new PyRecursiveElementVisitor() {
@Override
public void visitPyYieldExpression(PyYieldExpression node) {
final PyType type = context.getType(node);
if (node.isDelegating() && type instanceof PyCollectionType) {
final PyCollectionType collectionType = (PyCollectionType) type;
types.add(collectionType.getElementType(context));
} else {
types.add(type);
}
}
@Override
public void visitPyFunction(PyFunction node) {
// Ignore nested functions
}
});
final int n = types.size();
if (n == 1) {
elementType = Ref.create(types.iterator().next());
} else if (n > 0) {
elementType = Ref.create(PyUnionType.union(types));
}
}
if (elementType != null) {
final PyClass generator = cache.getClass(PyNames.FAKE_GENERATOR);
if (generator != null) {
return Ref.create(new PyCollectionTypeImpl(generator, false, elementType.get()));
}
}
if (!types.isEmpty()) {
return Ref.create(null);
}
return null;
}