diff options
Diffstat (limited to 'python/src/com/jetbrains/python/codeInsight/intentions/ImportToImportFromIntention.java')
-rw-r--r-- | python/src/com/jetbrains/python/codeInsight/intentions/ImportToImportFromIntention.java | 122 |
1 files changed, 69 insertions, 53 deletions
diff --git a/python/src/com/jetbrains/python/codeInsight/intentions/ImportToImportFromIntention.java b/python/src/com/jetbrains/python/codeInsight/intentions/ImportToImportFromIntention.java index fc35833473f8..e76240cefb1f 100644 --- a/python/src/com/jetbrains/python/codeInsight/intentions/ImportToImportFromIntention.java +++ b/python/src/com/jetbrains/python/codeInsight/intentions/ImportToImportFromIntention.java @@ -57,10 +57,11 @@ public class ImportToImportFromIntention implements IntentionAction { private PsiElement myReferee = null; private PyImportElement myImportElement = null; private Collection<PsiReference> myReferences = null; - private boolean myHasModuleReference = false; // is anything that resolves to our imported module is just an exact reference to that module + private boolean myHasModuleReference = false; + // is anything that resolves to our imported module is just an exact reference to that module private int myRelativeLevel; // true if "from ... import" - public IntentionState(Editor editor, PsiFile file) { + public IntentionState(@NotNull Editor editor, @NotNull PsiFile file) { boolean available = false; myImportElement = findImportElement(editor, file); if (myImportElement != null) { @@ -70,10 +71,10 @@ public class ImportToImportFromIntention implements IntentionAction { available = true; } else if (parent instanceof PyFromImportStatement) { - PyFromImportStatement from_import = (PyFromImportStatement)parent; - final int relative_level = from_import.getRelativeLevel(); - if (from_import.isValid() && relative_level > 0 && from_import.getImportSource() == null) { - myRelativeLevel = relative_level; + final PyFromImportStatement fromImport = (PyFromImportStatement)parent; + final int relativeLevel = fromImport.getRelativeLevel(); + if (fromImport.isValid() && relativeLevel > 0 && fromImport.getImportSource() == null) { + myRelativeLevel = relativeLevel; available = true; } } @@ -92,25 +93,27 @@ public class ImportToImportFromIntention implements IntentionAction { assert myImportElement != null : "isAvailable() must have returned true, but myImportElement is null"; // usages of imported name are qualifiers; what they refer to? - PyReferenceExpression reference = myImportElement.getImportReferenceExpression(); - if (reference != null) { - myModuleName = PyPsiUtils.toPath(reference); + final PyReferenceExpression importReference = myImportElement.getImportReferenceExpression(); + if (importReference != null) { + myModuleName = PyPsiUtils.toPath(importReference); myQualifierName = myImportElement.getVisibleName(); - myReferee = reference.getReference().resolve(); + myReferee = importReference.getReference().resolve(); myHasModuleReference = false; if (myReferee != null && myModuleName != null && myQualifierName != null) { final Collection<PsiReference> references = new ArrayList<PsiReference>(); PsiTreeUtil.processElements(file, new PsiElementProcessor() { public boolean execute(@NotNull PsiElement element) { if (element instanceof PyReferenceExpression && PsiTreeUtil.getParentOfType(element, PyImportElement.class) == null) { - PyReferenceExpression ref = (PyReferenceExpression)element; + final PyReferenceExpression ref = (PyReferenceExpression)element; if (myQualifierName.equals(PyPsiUtils.toPath(ref))) { // filter out other names that might resolve to our target - PsiElement parent_elt = ref.getParent(); - if (parent_elt instanceof PyQualifiedExpression) { // really qualified by us, not just referencing? - PsiElement resolved = ref.getReference().resolve(); + final PsiElement parentElt = ref.getParent(); + if (parentElt instanceof PyQualifiedExpression) { // really qualified by us, not just referencing? + final PsiElement resolved = ref.getReference().resolve(); if (resolved == myReferee) references.add(ref.getReference()); } - else myHasModuleReference = true; + else { + myHasModuleReference = true; + } } } return true; @@ -123,68 +126,70 @@ public class ImportToImportFromIntention implements IntentionAction { public void invoke() { assert myImportElement != null : "isAvailable() must have returned true, but myImportElement is null"; - PyUtil.sure(myImportElement.getImportReferenceExpression()); - Project project = myImportElement.getProject(); + sure(myImportElement.getImportReferenceExpression()); + final Project project = myImportElement.getProject(); + + final PyElementGenerator generator = PyElementGenerator.getInstance(project); + final LanguageLevel languageLevel = LanguageLevel.forElement(myImportElement); // usages of imported name are qualifiers; what they refer to? try { // remember names and make them drop qualifiers - Set<String> used_names = new HashSet<String>(); + final Set<String> usedNames = new HashSet<String>(); for (PsiReference ref : myReferences) { - PsiElement elt = ref.getElement(); - PsiElement parent_elt = elt.getParent(); - used_names.add(sure(PyUtil.sure(parent_elt).getLastChild()).getText()); // TODO: find ident node more properly + final PsiElement elt = ref.getElement(); + final PsiElement parentElt = elt.getParent(); + // TODO: find ident node more properly + final String nameUsed = sure(sure(parentElt).getLastChild()).getText(); + usedNames.add(nameUsed); if (!FileModificationService.getInstance().preparePsiElementForWrite(elt)) { return; } - PsiElement next_elt = elt.getNextSibling(); - if (next_elt != null && ".".equals(next_elt.getText())) next_elt.delete(); - elt.delete(); + assert parentElt instanceof PyReferenceExpression; + final PyElement newReference = generator.createExpressionFromText(languageLevel, nameUsed); + parentElt.replace(newReference); } // create a separate import stmt for the module - PsiElement importer = myImportElement.getParent(); - PyStatement import_statement; - PyImportElement[] import_elements; + final PsiElement importer = myImportElement.getParent(); + final PyStatement importStatement; + final PyImportElement[] importElements; if (importer instanceof PyImportStatement) { - import_statement = (PyImportStatement)importer; - import_elements = ((PyImportStatement)import_statement).getImportElements(); + importStatement = (PyImportStatement)importer; + importElements = ((PyImportStatement)importStatement).getImportElements(); } else if (importer instanceof PyFromImportStatement) { - import_statement = (PyFromImportStatement)importer; - import_elements = ((PyFromImportStatement)import_statement).getImportElements(); + importStatement = (PyFromImportStatement)importer; + importElements = ((PyFromImportStatement)importStatement).getImportElements(); } else { throw new IncorrectOperationException("Not an import at all"); } - PyElementGenerator generator = PyElementGenerator.getInstance(project); - StringBuilder builder = new StringBuilder("from ").append(getDots()).append(myModuleName).append(" import "); - builder.append(StringUtil.join(used_names, ", ")); - PyFromImportStatement from_import_stmt = - generator.createFromText(LanguageLevel.getDefault(), PyFromImportStatement.class, builder.toString()); - PsiElement parent = import_statement.getParent(); + final PyFromImportStatement newImportStatement = + generator.createFromImportStatement(languageLevel, getDots() + myModuleName, StringUtil.join(usedNames, ", "), null); + final PsiElement parent = importStatement.getParent(); sure(parent); sure(parent.isValid()); - if (import_elements.length == 1) { + if (importElements.length == 1) { if (myHasModuleReference) { - parent.addAfter(from_import_stmt, import_statement); // add 'import from': we need the module imported as is + parent.addAfter(newImportStatement, importStatement); // add 'import from': we need the module imported as is } else { // replace entire existing import - sure(parent.getNode()).replaceChild(sure(import_statement.getNode()), sure(from_import_stmt.getNode())); + sure(parent.getNode()).replaceChild(sure(importStatement.getNode()), sure(newImportStatement.getNode())); // import_statement.replace(from_import_stmt); } } else { if (!myHasModuleReference) { // cut the module out of import, add a from-import. - for (PyImportElement pie : import_elements) { + for (PyImportElement pie : importElements) { if (pie == myImportElement) { PyUtil.removeListNode(pie); break; } } } - parent.addAfter(from_import_stmt, import_statement); + parent.addAfter(newImportStatement, importStatement); } } catch (IncorrectOperationException ignored) { @@ -193,18 +198,24 @@ public class ImportToImportFromIntention implements IntentionAction { } + @NotNull public String getText() { - String module_name = "?"; + String moduleName = "?"; if (myImportElement != null) { - PyReferenceExpression reference = myImportElement.getImportReferenceExpression(); - if (reference != null) module_name = PyPsiUtils.toPath(reference); + final PyReferenceExpression reference = myImportElement.getImportReferenceExpression(); + if (reference != null) { + moduleName = PyPsiUtils.toPath(reference); + } } - return PyBundle.message("INTN.convert.to.from.$0.import.$1", getDots()+module_name, "..."); + return PyBundle.message("INTN.convert.to.from.$0.import.$1", getDots() + moduleName, "..."); } + @NotNull private String getDots() { String dots = ""; - for (int i=0; i<myRelativeLevel; i+=1) dots += "."; // this generally runs 1-2 times, so it's cheaper than allocating a StringBuilder + for (int i = 0; i < myRelativeLevel; i += 1) { + dots += "."; // this generally runs 1-2 times, so it's cheaper than allocating a StringBuilder + } return dots; } } @@ -222,10 +233,15 @@ public class ImportToImportFromIntention implements IntentionAction { } @Nullable - private static PyImportElement findImportElement(Editor editor, PsiFile file) { - PyImportElement import_elt = PsiTreeUtil.getParentOfType(file.findElementAt(editor.getCaretModel().getOffset()), PyImportElement.class); - if (import_elt != null && import_elt.isValid()) return import_elt; - else return null; + private static PyImportElement findImportElement(@NotNull Editor editor, @NotNull PsiFile file) { + final PsiElement elementAtCaret = file.findElementAt(editor.getCaretModel().getOffset()); + final PyImportElement importElement = PsiTreeUtil.getParentOfType(elementAtCaret, PyImportElement.class); + if (importElement != null && importElement.isValid()) { + return importElement; + } + else { + return null; + } } public boolean isAvailable(@NotNull Project project, Editor editor, PsiFile file) { @@ -233,7 +249,7 @@ public class ImportToImportFromIntention implements IntentionAction { return false; } - IntentionState state = new IntentionState(editor, file); + final IntentionState state = new IntentionState(editor, file); if (state.isAvailable()) { myText = state.getText(); return true; @@ -242,7 +258,7 @@ public class ImportToImportFromIntention implements IntentionAction { } public void invoke(@NotNull Project project, Editor editor, PsiFile file) throws IncorrectOperationException { - IntentionState state = new IntentionState(editor, file); + final IntentionState state = new IntentionState(editor, file); state.invoke(); } |