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(); }
@NotNull public String getPresentableText(@NotNull String myName) { final StringBuilder sb = new StringBuilder(getQualifiedName(myName, myPath, myImportElement)); PsiElement parent = null; if (myImportElement != null) { parent = myImportElement.getParent(); } if (myImportable instanceof PyFunction) { sb.append(((PyFunction) myImportable).getParameterList().getPresentableText(false)); } else if (myImportable instanceof PyClass) { final List<String> supers = ContainerUtil.mapNotNull( ((PyClass) myImportable).getSuperClasses(), new Function<PyClass, String>() { @Override public String fun(PyClass cls) { return PyUtil.isObjectClass(cls) ? null : cls.getName(); } }); if (!supers.isEmpty()) { sb.append("("); StringUtil.join(supers, ", ", sb); sb.append(")"); } } if (parent instanceof PyFromImportStatement) { sb.append(" from "); final PyFromImportStatement fromImportStatement = (PyFromImportStatement) parent; sb.append(StringUtil.repeat(".", fromImportStatement.getRelativeLevel())); final PyReferenceExpression source = fromImportStatement.getImportSource(); if (source != null) { sb.append(source.getReferencedName()); } } return sb.toString(); }