@Nullable private PyType resolveQualifierType( @NotNull List<Token<PyElementType>> tokens, @NotNull PyFile file, @NotNull TypeEvalContext context, @NotNull Map<TextRange, PyType> types, @NotNull Map<PyType, TextRange> fullRanges, @NotNull Map<PyType, PyImportElement> imports) { if (tokens.isEmpty()) { return null; } final Token<PyElementType> firstToken = tokens.get(0); final String firstText = firstToken.getText().toString(); final TextRange firstRange = firstToken.getRange(); final List<RatedResolveResult> resolveResults = file.multiResolveName(firstText); if (resolveResults.isEmpty()) { return getImplicitlyResolvedType(tokens, context, types, fullRanges, firstRange); } final List<PyType> members = Lists.newArrayList(); for (RatedResolveResult result : resolveResults) { final PsiElement resolved = result.getElement(); PyType type = null; if (resolved instanceof PyTargetExpression) { type = PyTypingTypeProvider.getTypeFromTargetExpression( (PyTargetExpression) resolved, context); } if (type == null && resolved instanceof PyTypedElement) { type = context.getType((PyTypedElement) resolved); } if (type != null) { if (!allowResolveToType(type)) { continue; } if (type instanceof PyClassLikeType) { type = ((PyClassLikeType) type).toInstance(); } types.put(firstRange, type); fullRanges.put(type, firstRange); for (PyFromImportStatement fromImportStatement : file.getFromImports()) { for (PyImportElement importElement : fromImportStatement.getImportElements()) { if (firstText.equals(importElement.getVisibleName())) { imports.put(type, importElement); } } } for (PyImportElement importElement : file.getImportTargets()) { if (firstText.equals(importElement.getVisibleName())) { imports.put(type, importElement); } } } members.add(type); } if (!members.isEmpty()) { tokens.remove(0); } return PyUnionType.union(members); }
private void addToExistingImport(PyImportElement src) { final PyElementGenerator gen = PyElementGenerator.getInstance(myTarget.getProject()); // did user choose 'import' or 'from import'? PsiElement parent = src.getParent(); if (parent instanceof PyFromImportStatement) { // add another import element right after the one we got PsiElement newImportElement = gen.createImportElement(LanguageLevel.getDefault(), myName); parent.add(newImportElement); } else { // just 'import' // all we need is to qualify our target myTarget.replace( gen.createExpressionFromText( LanguageLevel.forElement(myTarget), src.getVisibleName() + "." + myName)); } }
/** * Helper method that builds an import path, handling all these "import foo", "import foo as bar", * "from bar import foo", etc. Either importPath or importSource must be not null. * * @param name what is ultimately imported. * @param importPath known path to import the name. * @param source known ImportElement to import the name; its 'as' clause is used if present. * @return a properly qualified name. */ @NotNull public static String getQualifiedName( @NotNull String name, @Nullable QualifiedName importPath, @Nullable PyImportElement source) { final StringBuilder sb = new StringBuilder(); if (source != null) { final PsiElement parent = source.getParent(); if (parent instanceof PyFromImportStatement) { sb.append(name); } else { sb.append(source.getVisibleName()).append(".").append(name); } } else { if (importPath != null && importPath.getComponentCount() > 0) { sb.append(importPath).append("."); } sb.append(name); } return sb.toString(); }