PY-16221 Impl moving add from future imports when move

(cherry picked from commit b1cd64885d96c560976a29db9f5625a29df6d6ef)

IJ-MR-14916

GitOrigin-RevId: 5439442556b3d8ea2680d833109ac86559b765c1
This commit is contained in:
andrey.matveev
2021-09-20 16:52:32 +07:00
committed by intellij-monorepo-bot
parent e57fb63608
commit f4e98200ee
15 changed files with 100 additions and 3 deletions

View File

@@ -15,6 +15,7 @@ import com.intellij.util.ArrayUtil;
import com.intellij.util.IncorrectOperationException;
import com.intellij.util.containers.ContainerUtil;
import com.jetbrains.python.PyNames;
import com.jetbrains.python.codeInsight.imports.AddImportHelper;
import com.jetbrains.python.codeInsight.imports.PyImportOptimizer;
import com.jetbrains.python.psi.*;
import com.jetbrains.python.psi.impl.PyImportedModule;
@@ -24,9 +25,7 @@ import org.jetbrains.annotations.NonNls;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.*;
/**
* @author Dennis.Ushakov
@@ -37,6 +36,7 @@ public final class PyClassRefactoringUtil {
private static final Key<Boolean> ENCODED_USE_FROM_IMPORT = Key.create("PyEncodedUseFromImport");
private static final Key<String> ENCODED_IMPORT_AS = Key.create("PyEncodedImportAs");
private static final Key<List<PyReferenceExpression>> INJECTION_REFERENCES = Key.create("PyInjectionReferences");
private static final Key<Set<FutureFeature>> ENCODED_FROM_FUTURE_IMPORTS = Key.create("PyFromFutureImports");
private PyClassRefactoringUtil() {
@@ -185,6 +185,16 @@ public final class PyClassRefactoringUtil {
public static void restoreNamedReferences(@NotNull final PsiElement newElement,
@Nullable final PsiElement oldElement,
final PsiElement @NotNull [] otherMovedElements) {
Set<FutureFeature> fromFutureImports = newElement.getCopyableUserData(ENCODED_FROM_FUTURE_IMPORTS);
newElement.putCopyableUserData(ENCODED_FROM_FUTURE_IMPORTS, null);
PsiFile destFile = newElement.getContainingFile();
if (fromFutureImports != null & destFile != null) {
for (FutureFeature futureFeature: fromFutureImports) {
AddImportHelper.addOrUpdateFromImportStatement(destFile, PyNames.FUTURE_MODULE, futureFeature.toString(), null,
AddImportHelper.ImportPriority.FUTURE, null);
}
}
newElement.acceptChildren(new PyRecursiveElementVisitor() {
@Override
public void visitPyReferenceExpression(@NotNull PyReferenceExpression node) {
@@ -279,6 +289,12 @@ public final class PyClassRefactoringUtil {
* @param namesToSkip if reference inside of element has one of this names, it will not be saved.
*/
public static void rememberNamedReferences(@NotNull final PsiElement element, final String @NotNull ... namesToSkip) {
PsiFile containingFile = element.getContainingFile();
if (containingFile instanceof PyFile) {
Set<FutureFeature> fromFutureImports = collectFromFutureImports((PyFile)containingFile);
element.putCopyableUserData(ENCODED_FROM_FUTURE_IMPORTS, fromFutureImports);
}
element.accept(new PyRecursiveElementVisitor() {
@Override
public void visitPyReferenceExpression(@NotNull PyReferenceExpression node) {
@@ -339,6 +355,16 @@ public final class PyClassRefactoringUtil {
host.putCopyableUserData(INJECTION_REFERENCES, rememberedReferences);
}
private static @NotNull Set<FutureFeature> collectFromFutureImports(@NotNull PyFile file) {
EnumSet<FutureFeature> result = EnumSet.noneOf(FutureFeature.class);
for (FutureFeature feature: FutureFeature.values()) {
if (file.hasImportFromFuture(feature)) {
result.add(feature);
}
}
return result;
}
private static void rememberReference(@NotNull PyReferenceExpression node, @NotNull PsiElement element) {
// We will remember reference in deepest node (except for references to PyImportedModules, as we need references to modules, not to
// their packages)

View File

@@ -0,0 +1,6 @@
from __future__ import absolute_import, annotations
class NewParent(object):
def foo(self):
pass

View File

@@ -0,0 +1 @@
from __future__ import absolute_import

View File

@@ -0,0 +1,8 @@
from __future__ import annotations
from __future__ import absolute_import
from dest_module import NewParent
class MyClass(NewParent):
pass

View File

@@ -0,0 +1,6 @@
from __future__ import annotations
from __future__ import absolute_import
class MyClass(object):
def foo(self):
pass

View File

@@ -0,0 +1,5 @@
from __future__ import annotations
from __future__ import unicode_literals
from __future__ import division

View File

@@ -0,0 +1,6 @@
from __future__ import unicode_literals, division, annotations
from __future__ import print_function
class C:
pass

View File

@@ -0,0 +1,7 @@
from __future__ import annotations
from __future__ import unicode_literals
from __future__ import division
class C:
pass

View File

@@ -0,0 +1,2 @@
from __future__ import unicode_literals
from __future__ import print_function

View File

@@ -0,0 +1,4 @@
from __future__ import annotations
from __future__ import unicode_literals

View File

@@ -0,0 +1,5 @@
from __future__ import unicode_literals, annotations
class C:
pass

View File

@@ -0,0 +1,6 @@
from __future__ import annotations
from __future__ import unicode_literals
class C:
pass

View File

@@ -438,6 +438,16 @@ public class PyMoveTest extends PyTestCase {
doMoveSymbolTest("func", "dst.py");
}
// PY-16221
public void testFromFutureImports() {
doMoveSymbolTest("C", "b.py");
}
// PY-16221
public void testExistingFromFutureImportsNotDuplicated() {
doMoveSymbolTest("C", "b.py");
}
// PY-23831
public void testWithImportedForwardReferencesInTypeHints() {
doMoveSymbolTest("test", "dst.py");

View File

@@ -249,4 +249,9 @@ public final class PyExtractSuperclassTest extends PyClassRefactoringTest {
public void testNoClassCastExceptionInCopiedFunctionWithClassInitAndMethodCall() {
doSimpleTest("Baz", "Bar", null, true, false, ".baz");
}
// PY-16221
public void testFromFutureImports() {
multiFileTestHelper(".foo", false);
}
}