PY-54168 Moving class from __init__.py into own module removes it from __all__ attribute

GitOrigin-RevId: 534dfcb5ea6a6289d622236fed7cc859bc5ff030

python/src/com/jetbrains/python/refactoring/move/moduleMembers/PyMoveSymbolProcessor.java

@@ -2,6 +2,8 @@
 package com.jetbrains.python.refactoring.move.moduleMembers;
 
 import com.intellij.application.options.CodeStyle;
+import com.intellij.openapi.vfs.VfsUtilCore;
+import com.intellij.openapi.vfs.VirtualFile;
 import com.intellij.psi.*;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.psi.util.PsiUtilCore;
@@ -61,7 +63,8 @@ public class PyMoveSymbolProcessor {
           updateSingleUsage(usageElement, newElement);
         }
       }
-      final PsiElement[] unwrappedElements = ContainerUtil.mapNotNull(myAllMovedElements, SmartPsiElementPointer::getElement).toArray(PsiElement.EMPTY_ARRAY);
+      final PsiElement[] unwrappedElements =
+        ContainerUtil.mapNotNull(myAllMovedElements, SmartPsiElementPointer::getElement).toArray(PsiElement.EMPTY_ARRAY);
       PyClassRefactoringUtil.restoreNamedReferences(newElementBody, myMovedElement, unwrappedElements);
       final PyTrailingBlankLinesPostFormatProcessor postFormatProcessor = new PyTrailingBlankLinesPostFormatProcessor();
       if (PsiTreeUtil.nextVisibleLeaf(newElementBody) == null) {
@@ -129,7 +132,32 @@ public class PyMoveSymbolProcessor {
     else if (usage instanceof PyStringLiteralExpression) {
       for (PsiReference ref : usage.getReferences()) {
         if (ref instanceof PyDunderAllReference) {
-          usage.delete();
+          if (myMovedElement.getContainingFile().getName().equals(PyNames.INIT_DOT_PY)) {
+            // Update import statement in `__init__.py` if symbol is moved from there to another module
+            VirtualFile targetFile = newElement.getContainingFile().getVirtualFile();
+            VirtualFile currentFile = usage.getContainingFile().getVirtualFile();
+            if (targetFile != null && currentFile != null) {
+              VirtualFile currentDir = currentFile.getParent();
+              if (currentDir != null && !VfsUtilCore.isAncestor(currentDir, targetFile, true)) {
+                usage.delete();
+              }
+              PyPsiRefactoringUtil.insertImport(usage, newElement);
+            }
+          }
+          else if (ref.getElement().getContainingFile().getName().equals(PyNames.INIT_DOT_PY)) {
+            // Remove old import statement in `__init__.py` if symbol is moved from some ancestor module to another one
+            VirtualFile targetFile = newElement.getContainingFile().getVirtualFile();
+            VirtualFile currentFile = usage.getContainingFile().getVirtualFile();
+            if (targetFile != null && currentFile != null) {
+              VirtualFile currentDir = currentFile.getParent();
+              if (currentDir != null && !VfsUtilCore.isAncestor(currentDir, targetFile, true)) {
+                usage.delete();
+              }
+            }
+          }
+          else {
+            usage.delete();
+          }
         }
         else {
           if (ref.isReferenceTo(myMovedElement)) {
@@ -204,7 +232,7 @@ public class PyMoveSymbolProcessor {
       }
     }
     else {
-      final PsiReference ref = usage.getReference();  
+      final PsiReference ref = usage.getReference();
       if (ref != null) {
         resolvedElements.add(ref.resolve());
       }

python/testData/refactoring/move/moveFromInitPyPreserveDunderAll/after/src/__init__.py

@@ -0,0 +1,5 @@
+__all__ = ['MyClass']
+
+from my_class import MyClass
+
+

python/testData/refactoring/move/moveFromInitPyPreserveDunderAll/after/src/my_class.py

@@ -0,0 +1,3 @@
+class MyClass:
+    def __init__(self):
+        pass

python/testData/refactoring/move/moveFromInitPyPreserveDunderAll/before/src/__init__.py

@@ -0,0 +1,6 @@
+__all__ = ['MyClass']
+
+
+class MyClass:
+    def __init__(self):
+        pass

python/testData/refactoring/move/moveInHierarchyUpdatesDunderAllInInitPy/after/src/__init__.py

@@ -0,0 +1,3 @@
+__all__ = ['MyClass']
+
+from new import MyClass

python/testData/refactoring/move/moveInHierarchyUpdatesDunderAllInInitPy/after/src/new.py

@@ -0,0 +1,3 @@
+class MyClass:
+    def __init__(self):
+        pass

python/testData/refactoring/move/moveInHierarchyUpdatesDunderAllInInitPy/before/src/__init__.py

@@ -0,0 +1,4 @@
+__all__ = ['MyClass']
+
+
+from my_class import MyClass

python/testData/refactoring/move/moveInHierarchyUpdatesDunderAllInInitPy/before/src/my_class.py

@@ -0,0 +1,3 @@
+class MyClass:
+    def __init__(self):
+        pass

python/testSrc/com/jetbrains/python/refactoring/PyMoveTest.java

@@ -312,6 +312,16 @@ public class PyMoveTest extends PyTestCase {
     doMoveSymbolTest("func", "b.py");
   }
 
+  // PY-54168
+  public void testMoveFromInitPyPreserveDunderAll() {
+    doMoveSymbolTest("MyClass", "my_class.py");
+  }
+
+  // PY-54168
+  public void testMoveInHierarchyUpdatesDunderAllInInitPy() {
+    doMoveSymbolTest("MyClass", "new.py");
+  }
+
   // PY-15343
   public void testDunderAllSingleElementTuple() {
     doMoveSymbolTest("func", "b.py");