@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;
}
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 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;
}
protected static PsiElement findAnchor(List<PsiElement> occurrences) {
PsiElement anchor = occurrences.get(0);
next:
do {
final PyStatement statement = PsiTreeUtil.getParentOfType(anchor, PyStatement.class);
if (statement != null) {
final PsiElement parent = statement.getParent();
for (PsiElement element : occurrences) {
if (!PsiTreeUtil.isAncestor(parent, element, true)) {
anchor = statement;
continue next;
}
}
}
return statement;
} while (true);
}
private boolean smartIntroduce(final IntroduceOperation operation) {
final Editor editor = operation.getEditor();
final PsiFile file = operation.getFile();
int offset = editor.getCaretModel().getOffset();
PsiElement elementAtCaret = file.findElementAt(offset);
if (!checkIntroduceContext(file, editor, elementAtCaret)) return true;
final List<PyExpression> expressions = new ArrayList<PyExpression>();
while (elementAtCaret != null) {
if (elementAtCaret instanceof PyStatement || elementAtCaret instanceof PyFile) {
break;
}
if (elementAtCaret instanceof PyExpression && isValidIntroduceVariant(elementAtCaret)) {
expressions.add((PyExpression) elementAtCaret);
}
elementAtCaret = elementAtCaret.getParent();
}
if (expressions.size() == 1 || ApplicationManager.getApplication().isUnitTestMode()) {
operation.setElement(expressions.get(0));
performActionOnElement(operation);
return true;
} else if (expressions.size() > 1) {
IntroduceTargetChooser.showChooser(
editor,
expressions,
new Pass<PyExpression>() {
@Override
public void pass(PyExpression pyExpression) {
operation.setElement(pyExpression);
performActionOnElement(operation);
}
},
new Function<PyExpression, String>() {
public String fun(PyExpression pyExpression) {
return pyExpression.getText();
}
});
return true;
}
return false;
}
@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;
}
@Nullable
@Override
public ParseResult fun(@NotNull Pair<Token<PyElementType>, List<Token<PyElementType>>> value) {
final Token<PyElementType> first = value.getFirst();
final List<Token<PyElementType>> rest = value.getSecond();
final TextRange firstRange = first.getRange();
final boolean unqualified = rest.isEmpty();
if (unqualified) {
final ParseResult result = parseBuiltinType(first);
if (result != null) {
return result;
}
}
final PsiFile file = myAnchor.getContainingFile();
final List<Token<PyElementType>> tokens = new ArrayList<>();
tokens.add(first);
tokens.addAll(rest);
if (file instanceof PyFile) {
final PyFile pyFile = (PyFile) file;
final TypeEvalContext context = TypeEvalContext.codeInsightFallback(file.getProject());
final Map<TextRange, PyType> types = new HashMap<>();
final Map<PyType, TextRange> fullRanges = new HashMap<>();
final Map<PyType, PyImportElement> imports = new HashMap<>();
PyType type = resolveQualifierType(tokens, pyFile, context, types, fullRanges, imports);
PsiElement resolved = type != null ? getElement(type) : null;
if (type != null) {
final PyResolveContext resolveContext =
PyResolveContext.defaultContext().withTypeEvalContext(context);
final PyExpression expression =
myAnchor instanceof PyExpression ? (PyExpression) myAnchor : null;
for (Token<PyElementType> token : tokens) {
final PyType qualifierType = type;
type = null;
final List<? extends RatedResolveResult> results =
qualifierType.resolveMember(
token.getText().toString(), expression, AccessDirection.READ, resolveContext);
if (results != null && !results.isEmpty()) {
resolved = results.get(0).getElement();
if (resolved instanceof PyTypedElement) {
type = context.getType((PyTypedElement) resolved);
if (type != null && !allowResolveToType(type)) {
type = null;
break;
}
if (type instanceof PyClassLikeType) {
type = ((PyClassLikeType) type).toInstance();
}
}
}
if (type == null) {
break;
}
types.put(token.getRange(), type);
fullRanges.put(
type,
TextRange.create(firstRange.getStartOffset(), token.getRange().getEndOffset()));
}
if (type != null) {
return new ParseResult(resolved, type, types, fullRanges, imports);
}
}
}
return EMPTY_RESULT;
}
