@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 String extractReturnType() { final String ARROW = "->"; final StructuredDocString structuredDocString = getStructuredDocString(); if (structuredDocString != null) { return structuredDocString.getReturnType(); } final String docString = getDocStringValue(); if (docString != null && docString.contains(ARROW)) { final List<String> lines = StringUtil.split(docString, "\n"); while (lines.size() > 0 && lines.get(0).trim().length() == 0) { lines.remove(0); } if (lines.size() > 1 && lines.get(1).trim().length() == 0) { String firstLine = lines.get(0); int pos = firstLine.lastIndexOf(ARROW); if (pos >= 0) { return firstLine.substring(pos + 2).trim(); } } } 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; }
private static PyType getTypeByControlFlow( @NotNull String name, @NotNull TypeEvalContext context, @NotNull PyExpression anchor, @NotNull ScopeOwner scopeOwner) { PyAugAssignmentStatement augAssignment = PsiTreeUtil.getParentOfType(anchor, PyAugAssignmentStatement.class); try { final PyElement element = augAssignment != null ? augAssignment : anchor; final List<ReadWriteInstruction> defs = PyDefUseUtil.getLatestDefs(scopeOwner, name, element, true); if (!defs.isEmpty()) { PyType type = defs.get(0).getType(context, anchor); for (int i = 1; i < defs.size(); i++) { type = PyUnionType.union(type, defs.get(i).getType(context, anchor)); } return type; } } catch (PyDefUseUtil.InstructionNotFoundException ignored) { } return null; }
@Nullable private Modifier getWrappersFromStub() { final StubElement parentStub = getStub().getParentStub(); final List childrenStubs = parentStub.getChildrenStubs(); int index = childrenStubs.indexOf(getStub()); if (index >= 0 && index < childrenStubs.size() - 1) { StubElement nextStub = (StubElement) childrenStubs.get(index + 1); if (nextStub instanceof PyTargetExpressionStub) { final PyTargetExpressionStub targetExpressionStub = (PyTargetExpressionStub) nextStub; if (targetExpressionStub.getInitializerType() == PyTargetExpressionStub.InitializerType.CallExpression) { final QualifiedName qualifiedName = targetExpressionStub.getInitializer(); if (QualifiedName.fromComponents(PyNames.CLASSMETHOD).equals(qualifiedName)) { return CLASSMETHOD; } if (QualifiedName.fromComponents(PyNames.STATICMETHOD).equals(qualifiedName)) { return STATICMETHOD; } } } } return null; }