PY-23968 PyFromImportStatementImpl.addInternal supports adding new names "before" and at the beginning

GitOrigin-RevId: 838658cf749f62f0da1e9656021a800515073926
This commit is contained in:
Mikhail Golubev
2020-03-06 14:07:46 +00:00
committed by intellij-monorepo-bot
parent e5a732aec9
commit 5e814f0e78
12 changed files with 108 additions and 24 deletions
@@ -187,32 +187,41 @@ public class PyFromImportStatementImpl extends PyBaseElementImpl<PyFromImportSta
@Override
public ASTNode addInternal(ASTNode first, ASTNode last, ASTNode anchor, Boolean before) {
boolean addingNewName = first == last &&
first.getElementType() == PyElementTypes.IMPORT_ELEMENT &&
(anchor == null || anchor.getElementType() == PyElementTypes.IMPORT_ELEMENT);
if (!addingNewName) {
return super.addInternal(first, last, anchor, before);
}
if (anchor == null) {
// adding last element; the import may be "from ... import (...)", must get before the last ")"
PsiElement lastChild = getLastChild();
if (lastChild != null) {
while (lastChild instanceof PsiComment) {
lastChild = lastChild.getPrevSibling();
anchor = lastChild.getNode();
final PyImportElement[] elements = getImportElements();
if (elements.length != 0) {
if (before) {
anchor = elements[elements.length - 1].getNode();
before = false;
}
else {
anchor = elements[0].getNode();
before = true;
}
ASTNode rpar_node = lastChild.getNode();
if (rpar_node != null && rpar_node.getElementType() == PyTokenTypes.RPAR) anchor = rpar_node;
}
}
final ASTNode result = super.addInternal(first, last, anchor, before);
ASTNode prevNode = result;
do {
prevNode = prevNode.getTreePrev();
}
while (prevNode != null && prevNode.getElementType() == TokenType.WHITE_SPACE);
if (prevNode != null && prevNode.getElementType() == PyElementTypes.IMPORT_ELEMENT &&
result.getElementType() == PyElementTypes.IMPORT_ELEMENT) {
ASTNode comma = PyElementGenerator.getInstance(getProject()).createComma();
super.addInternal(comma, comma, prevNode, false);
// In an incomplete from import statement there is a special sentinel empty PyImportElement at the end
if (anchor != null && anchor.getTextLength() == 0) {
getNode().replaceChild(anchor, first);
return first;
}
else {
final ASTNode result = super.addInternal(first, last, anchor, before);
if (anchor != null && anchor.getElementType() == PyElementTypes.IMPORT_ELEMENT &&
result.getElementType() == PyElementTypes.IMPORT_ELEMENT) {
ASTNode comma = PyElementGenerator.getInstance(getProject()).createComma();
super.addInternal(comma, comma, before ? result : anchor, false);
}
return result;
}
return result;
}
@Override
@@ -0,0 +1,3 @@
from lib import b, c
print(<error descr="Unresolved reference 'a'"><caret>a</error>, b, c)
@@ -0,0 +1 @@
a, b, c = 1, 2, 3
@@ -0,0 +1,3 @@
from lib import a, b
print(a, b, <error descr="Unresolved reference 'c'"><caret>c</error>)
@@ -0,0 +1,3 @@
from lib import a, b, c
print(a, b, c)
@@ -0,0 +1,3 @@
from lib import a, b, c
print(a, b, c)
@@ -16,10 +16,9 @@
package com.jetbrains.python;
import com.intellij.psi.util.QualifiedName;
import com.intellij.util.containers.ContainerUtil;
import com.jetbrains.python.fixtures.PyTestCase;
import com.jetbrains.python.psi.LanguageLevel;
import com.jetbrains.python.psi.PyElementGenerator;
import com.jetbrains.python.psi.PyQualifiedExpression;
import com.jetbrains.python.psi.*;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
@@ -41,6 +40,54 @@ public class PyMiscellaneousPsiOperationsTest extends PyTestCase {
checkAsQualifiedNameResult("-foo + bar", null);
}
public void testAddingNameInFromImportStatement() {
checkAddingNameInFromImport("from mod import foo", "bar", "foo", true, "from mod import bar, foo");
checkAddingNameInFromImport("from mod import foo", "bar", "foo", false, "from mod import foo, bar");
checkAddingNameInFromImport("from mod import foo", "bar", null, false, "from mod import bar, foo");
checkAddingNameInFromImport("from mod import foo", "bar", null, true, "from mod import foo, bar");
checkAddingNameInFromImport("from mod import (foo) # comment", "bar", "foo", true, "from mod import (bar, foo) # comment");
checkAddingNameInFromImport("from mod import (foo) # comment", "bar", "foo", false, "from mod import (foo, bar) # comment");
checkAddingNameInFromImport("from mod import (foo,)", "bar", "foo", false, "from mod import (foo, bar,)");
checkAddingNameInFromImport("from mod import (foo # comment\n" +
" )",
"bar", "foo", false,
"from mod import (foo, bar # comment\n" +
" )");
checkAddingNameInFromImport("from mod import ", "bar", null, false, "from mod import bar");
checkAddingNameInFromImport("from mod import ", "bar", null, true, "from mod import bar");
checkAddingNameInFromImport("from mod import (", "bar", null, true, "from mod import (bar");
// TODO change where the placeholder empty import element is added in such cases
//checkAddingNameInFromImport("from mod import ( # comment", "bar", null, true, "from mod import (bar # comment");
checkAddingNameInFromImport("from mod import ()", "bar", null, true, "from mod import (bar)");
}
private void checkAddingNameInFromImport(@NotNull String fromImport,
@NotNull String newName,
@Nullable String anchorName,
boolean before,
@NotNull String result) {
final PyElementGenerator generator = PyElementGenerator.getInstance(myFixture.getProject());
final LanguageLevel languageLevel = LanguageLevel.PYTHON27;
final PyFromImportStatement fromImportElem = generator.createFromText(languageLevel, PyFromImportStatement.class, fromImport, new int[]{0});
final PyImportElement anchor;
if (anchorName != null) {
anchor = ContainerUtil.find(fromImportElem.getImportElements(),
importElem -> importElem.getImportedQName().toString().equals(anchorName));
assertNotNull(anchor);
}
else {
anchor = null;
}
final PyImportElement newNameElem = generator.createImportElement(languageLevel, newName, null);
if (before) {
fromImportElem.addBefore(newNameElem, anchor);
}
else {
fromImportElem.addAfter(newNameElem, anchor);
}
assertEquals(result, fromImportElem.getText());
}
private void checkAsQualifiedNameResult(@NotNull String expression, @Nullable QualifiedName expectedQualifiedName) {
final PyElementGenerator generator = PyElementGenerator.getInstance(myFixture.getProject());
final PyQualifiedExpression expr = (PyQualifiedExpression)generator.createExpressionFromText(LanguageLevel.PYTHON27, expression);
@@ -224,7 +224,21 @@ public class PyAddImportQuickFixTest extends PyQuickFixTestCase {
}
// PY-23968
public void testOrderingOfNamesInFromImport() {
public void testOrderingOfNamesInFromImportBeginning() {
getPythonCodeStyleSettings().OPTIMIZE_IMPORTS_SORT_IMPORTS = true;
getPythonCodeStyleSettings().OPTIMIZE_IMPORTS_SORT_NAMES_IN_FROM_IMPORTS = true;
doMultiFileAutoImportTest("Import");
}
// PY-23968
public void testOrderingOfNamesInFromImportInTheMiddle() {
getPythonCodeStyleSettings().OPTIMIZE_IMPORTS_SORT_IMPORTS = true;
getPythonCodeStyleSettings().OPTIMIZE_IMPORTS_SORT_NAMES_IN_FROM_IMPORTS = true;
doMultiFileAutoImportTest("Import");
}
// PY-23968
public void testOrderingOfNamesInFromImportEnd() {
getPythonCodeStyleSettings().OPTIMIZE_IMPORTS_SORT_IMPORTS = true;
getPythonCodeStyleSettings().OPTIMIZE_IMPORTS_SORT_NAMES_IN_FROM_IMPORTS = true;
doMultiFileAutoImportTest("Import");