@Nullable public static PyType getType(@NotNull PsiElement resolved, @NotNull List<PyType> elementTypes) { final String qualifiedName = getQualifiedName(resolved); final List<Integer> paramListTypePositions = new ArrayList<>(); final List<Integer> ellipsisTypePositions = new ArrayList<>(); for (int i = 0; i < elementTypes.size(); i++) { final PyType type = elementTypes.get(i); if (type instanceof PyTypeParser.ParameterListType) { paramListTypePositions.add(i); } else if (type instanceof PyTypeParser.EllipsisType) { ellipsisTypePositions.add(i); } } if (!paramListTypePositions.isEmpty()) { if (!("typing.Callable".equals(qualifiedName) && paramListTypePositions.equals(list(0)))) { return null; } } if (!ellipsisTypePositions.isEmpty()) { if (!("typing.Callable".equals(qualifiedName) && ellipsisTypePositions.equals(list(0)) || "typing.Tuple".equals(qualifiedName) && ellipsisTypePositions.equals(list(1)) && elementTypes.size() == 2)) { return null; } } if ("typing.Union".equals(qualifiedName)) { return PyUnionType.union(elementTypes); } if ("typing.Optional".equals(qualifiedName) && elementTypes.size() == 1) { return PyUnionType.union(elementTypes.get(0), PyNoneType.INSTANCE); } if ("typing.Callable".equals(qualifiedName) && elementTypes.size() == 2) { final PyTypeParser.ParameterListType paramList = as(elementTypes.get(0), PyTypeParser.ParameterListType.class); if (paramList != null) { return new PyCallableTypeImpl(paramList.getCallableParameters(), elementTypes.get(1)); } if (elementTypes.get(0) instanceof PyTypeParser.EllipsisType) { return new PyCallableTypeImpl(null, elementTypes.get(1)); } } if ("typing.Tuple".equals(qualifiedName)) { if (elementTypes.size() > 1 && elementTypes.get(1) instanceof PyTypeParser.EllipsisType) { return PyTupleType.createHomogeneous(resolved, elementTypes.get(0)); } return PyTupleType.create(resolved, elementTypes); } final PyType builtinCollection = getBuiltinCollection(resolved); if (builtinCollection instanceof PyClassType) { final PyClassType classType = (PyClassType) builtinCollection; return new PyCollectionTypeImpl(classType.getPyClass(), false, elementTypes); } return null; }
@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 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; }
@NotNull public static ParseResult parsePep484FunctionTypeComment( @NotNull PsiElement anchor, @NotNull String text) { final ForwardDeclaration<ParseResult, PyElementType> typeExpr = ForwardDeclaration.create(); final Function<Pair<ParseResult, List<ParseResult>>, ParseResult> toParamTypeList = pair -> { if (pair != null) { final ParseResult first = pair.getFirst(); final List<ParseResult> second = pair.getSecond(); final List<PyType> itemTypes = new ArrayList<>(); ParseResult result = first; itemTypes.add(result.getType()); for (ParseResult r : second) { result = result.merge(r); itemTypes.add(r.getType()); } return result.withType(ParameterListType.fromParameterTypes(itemTypes)); } return EMPTY_RESULT.withType(new ParameterListType(Collections.emptyList())); }; final FunctionalParser<ParseResult, PyElementType> ellipsis = op("...").map(token -> EMPTY_RESULT.withType(EllipsisType.INSTANCE)).named("ellipsis"); final FunctionalParser<ParseResult, PyElementType> classType = token(IDENTIFIER) .then(many(op(".").skipThen(token(IDENTIFIER)))) .map(new MakeSimpleType(anchor)) .cached() .named("class-type"); final FunctionalParser<ParseResult, PyElementType> typeParamList = op("[") .skipThen(maybe(typeExpr.then(many(op(",").skipThen(typeExpr))))) .thenSkip(op("]")) .map(toParamTypeList) .named("type-param-list"); final FunctionalParser<ParseResult, PyElementType> typeParam = typeExpr.or(typeParamList).or(ellipsis).named("type-param"); final FunctionalParser<ParseResult, PyElementType> genericType = classType .thenSkip(op("[")) .then(typeParam) .then(many(op(",").skipThen(typeParam))) .thenSkip(op("]")) .map( value -> { final Pair<ParseResult, ParseResult> firstPair = value.getFirst(); final ParseResult first = firstPair.getFirst(); final ParseResult second = firstPair.getSecond(); final List<ParseResult> third = value.getSecond(); final List<PyType> typesInBrackets = new ArrayList<>(); typesInBrackets.add(second.getType()); ParseResult result = first; result = result.merge(second); for (ParseResult r : third) { typesInBrackets.add(r.getType()); result = result.merge(r); } final List<PyType> elementTypes = third.isEmpty() ? Collections.singletonList(second.getType()) : typesInBrackets; final PsiElement resolved = first.getElement(); if (resolved != null) { final PyType typingType = PyTypingTypeProvider.getType(resolved, elementTypes); if (typingType != null) { return result.withType(typingType); } } return EMPTY_RESULT; }) .named("generic-type"); typeExpr.define(genericType.or(classType)).named("type-expr"); final FunctionalParser<ParseResult, PyElementType> paramType = maybe(op("*")) .then(maybe(op("*"))) .then(typeExpr) .map( pair -> { final ParseResult paramResult = pair.getSecond(); final PyType type = paramResult.getType(); int starCount = 0; if (pair.getFirst().getFirst() != null) { starCount++; } if (pair.getFirst().getSecond() != null) { starCount++; } if (starCount == 0) { return paramResult; } else if (starCount == 1) { final PyClassType tupleType = PyTupleType.createHomogeneous(anchor, type); if (tupleType != null) { return paramResult.withType(tupleType); } return EMPTY_RESULT; } else if (starCount == 2) { final PyBuiltinCache builtinCache = PyBuiltinCache.getInstance(anchor); final PyClassType dictType = builtinCache.getDictType(); if (dictType != null) { final PyClass pyClass = dictType.getPyClass(); return paramResult.withType( new PyCollectionTypeImpl( pyClass, false, Arrays.asList(builtinCache.getStrType(), type))); } return EMPTY_RESULT; } return EMPTY_RESULT; }) .named("param-type"); final FunctionalParser<ParseResult, PyElementType> paramTypes = paramType.then(many(op(",").skipThen(paramType))).map(toParamTypeList).named("param-types"); final FunctionalParser<ParseResult, PyElementType> funcType = op("(") .skipThen(maybe(paramTypes.or(ellipsis))) .thenSkip(op(")")) .thenSkip(op("->")) .then(typeExpr) .map( value -> { final ParseResult paramsResult = value.getFirst(); final ParseResult returnResult = value.getSecond(); final List<PyCallableParameter> parameters; ParseResult result = returnResult; if (paramsResult != null) { result = result.merge(paramsResult); final ParameterListType paramTypesList = as(paramsResult.getType(), ParameterListType.class); if (paramTypesList != null) { parameters = paramTypesList.getCallableParameters(); } // ellipsis else { parameters = null; } } else { parameters = Collections.emptyList(); result = returnResult; } return result.withType( new PyCallableTypeImpl(parameters, returnResult.getType())); }) .named("func-type"); final FunctionalParser<ParseResult, PyElementType> typeFile = funcType.endOfInput().named("function-type-comment"); try { return typeFile.parse(tokenize(text)); } catch (ParserException e) { return EMPTY_RESULT; } }
/** @param anchor should never be null or null will be returned */ @NotNull public static ParseResult parse(@Nullable final PsiElement anchor, @NotNull String type) { PyPsiUtils.assertValid(anchor); if (anchor == null) { return EMPTY_RESULT; } final ForwardDeclaration<ParseResult, PyElementType> typeExpr = ForwardDeclaration.create(); final FunctionalParser<ParseResult, PyElementType> classType = token(IDENTIFIER) .then(many(op(".").skipThen(token(IDENTIFIER)))) .map(new MakeSimpleType(anchor)) .cached() .named("class-type"); final FunctionalParser<ParseResult, PyElementType> tupleType = op("(") .skipThen(typeExpr) .then(many(op(",").skipThen(typeExpr))) .thenSkip(op(")")) .map( value -> { ParseResult result = value.getFirst(); final List<ParseResult> rest = value.getSecond(); if (rest.isEmpty()) { return result; } final List<PyType> types = new ArrayList<>(); types.add(result.getType()); for (ParseResult r : rest) { result = result.merge(r); types.add(r.getType()); } return result.withType(PyTupleType.create(anchor, types)); }) .named("tuple-type"); final FunctionalParser<ParseResult, PyElementType> typeParameter = token(PARAMETER) .then(maybe(op("<=").skipThen(typeExpr))) .map( value -> { final Token<PyElementType> token = value.getFirst(); final String name = token.getText().toString(); final TextRange range = token.getRange(); final ParseResult boundResult = value.getSecond(); if (boundResult != null) { final PyGenericType type1 = new PyGenericType(name, boundResult.getType()); final ParseResult result = new ParseResult(null, type1, range); return result.merge(boundResult).withType(type1); } return new ParseResult(null, new PyGenericType(name, null), range); }) .named("type-parameter"); final FunctionalParser<ParseResult, PyElementType> simpleExpr = classType.or(tupleType).or(typeParameter).named("simple-expr"); final FunctionalParser<ParseResult, PyElementType> paramExpr = classType .thenSkip(op("[")) .then(typeExpr) .then(many(op(",").skipThen(typeExpr))) .thenSkip(op("]")) .map( value -> { final Pair<ParseResult, ParseResult> firstPair = value.getFirst(); final ParseResult first = firstPair.getFirst(); final ParseResult second = firstPair.getSecond(); final List<ParseResult> third = value.getSecond(); final PyType firstType = first.getType(); final List<PyType> typesInBrackets = new ArrayList<>(); typesInBrackets.add(second.getType()); ParseResult result = first; result = result.merge(second); for (ParseResult r : third) { typesInBrackets.add(r.getType()); result = result.merge(r); } final List<PyType> elementTypes = third.isEmpty() ? Collections.singletonList(second.getType()) : typesInBrackets; final PsiElement resolved = first.getElement(); if (resolved != null) { final PyType typingType = PyTypingTypeProvider.getType(resolved, elementTypes); if (typingType != null) { return result.withType(typingType); } } if (firstType instanceof PyClassType) { final PyType type1 = new PyCollectionTypeImpl( ((PyClassType) firstType).getPyClass(), false, elementTypes); return result.withType(type1); } return EMPTY_RESULT; }) .or( classType .thenSkip(op("of")) .then(simpleExpr) .map( value -> { final ParseResult firstResult = value.getFirst(); final ParseResult secondResult = value.getSecond(); final ParseResult result = firstResult.merge(secondResult); final PyType firstType = firstResult.getType(); final PyType secondType = secondResult.getType(); if (firstType != null) { if (firstType instanceof PyClassType && secondType != null) { return result.withType( new PyCollectionTypeImpl( ((PyClassType) firstType).getPyClass(), false, Collections.singletonList(secondType))); } return result.withType(firstType); } return EMPTY_RESULT; })) .or( classType .thenSkip(op("from")) .then(simpleExpr) .thenSkip(op("to")) .then(simpleExpr) .map( value -> { final Pair<ParseResult, ParseResult> firstPair = value.getFirst(); final ParseResult first = firstPair.getFirst(); final ParseResult second = firstPair.getSecond(); final ParseResult third = value.getSecond(); final PyType firstType = first.getType(); if (firstType instanceof PyClassType) { final List<PyType> elementTypes = Arrays.asList(second.getType(), third.getType()); final PyCollectionTypeImpl type1 = new PyCollectionTypeImpl( ((PyClassType) firstType).getPyClass(), false, elementTypes); return first.merge(second).merge(third).withType(type1); } return EMPTY_RESULT; })) .named("param-expr"); final FunctionalParser<ParseResult, PyElementType> callableExpr = op("(") .skipThen(maybe(typeExpr.then(many(op(",").skipThen(typeExpr))))) .thenSkip(op(")")) .thenSkip(op("->")) .then(typeExpr) .map( value -> { final List<PyCallableParameter> parameters = new ArrayList<>(); final ParseResult returnResult = value.getSecond(); ParseResult result; final Pair<ParseResult, List<ParseResult>> firstPair = value.getFirst(); if (firstPair != null) { final ParseResult first = firstPair.getFirst(); final List<ParseResult> second = firstPair.getSecond(); result = first; parameters.add(new PyCallableParameterImpl(null, first.getType())); for (ParseResult r : second) { result = result.merge(r); parameters.add(new PyCallableParameterImpl(null, r.getType())); } result = result.merge(returnResult); } else { result = returnResult; } return result.withType( new PyCallableTypeImpl(parameters, returnResult.getType())); }) .named("callable-expr"); final FunctionalParser<ParseResult, PyElementType> singleExpr = paramExpr.or(callableExpr).or(simpleExpr).named("single-expr"); final FunctionalParser<ParseResult, PyElementType> unionExpr = singleExpr .then(many(op("or").or(op("|")).skipThen(singleExpr))) .map( value -> { final ParseResult first = value.getFirst(); final List<ParseResult> rest = value.getSecond(); if (rest.isEmpty()) { return first; } final List<PyType> types = new ArrayList<>(); types.add(first.getType()); ParseResult result = first; for (ParseResult r : rest) { types.add(r.getType()); result = result.merge(r); } return result.withType(PyUnionType.union(types)); }) .named("union-expr"); typeExpr.define(unionExpr).named("type-expr"); final FunctionalParser<ParseResult, PyElementType> typeFile = typeExpr.endOfInput().named("type-file"); try { return typeFile.parse(tokenize(type)); } catch (ParserException e) { return EMPTY_RESULT; } }