PY-14595 Handle namespace packages when updating imports for moved symbol

I added new method PyUtil#turnDirIntoPackageElement, that unlike
PyUtil#turnDirIntoInit is aware about namespace packages and return
passed directory itself in this case. Not sure, that altering
behavior of #turnDirIntoInit is a good idea, because its name will be
confusing then.

python/src/com/jetbrains/python/psi/PyUtil.java

@@ -967,6 +967,32 @@ public class PyUtil {
     } // don't touch non-dirs
   }
 
+  /**
+   * If target is a PsiDirectory, that is also a valid Python package, return PsiFile that points to __init__.py,
+   * if such file exists, or directory itself (i.e. namespace package). Otherwise, return {@code null}.
+   * Unlike {@link #turnDirIntoInit(com.intellij.psi.PsiElement)} this function handles namespace packages and
+   * accepts only PsiDirectories as target.
+   *
+   * @param target directory to check
+   * @param anchor optional PSI element to determine language level as for {@link #isPackage(com.intellij.psi.PsiDirectory, com.intellij.psi.PsiElement)}
+   * @return PsiFile or PsiDirectory, if target is a Python package and {@code null} null otherwise
+   */
+  @Nullable
+  public static PsiElement turnDirIntoPackageElement(@NotNull PsiDirectory target, @Nullable PsiElement anchor) {
+    if (isPackage(target, anchor)) {
+      final PsiFile file = target.findFile(PyNames.INIT_DOT_PY);
+      return file != null ? file : target;
+    }
+    else {
+      return null;
+    }
+  }
+
+  /**
+   * If target is a Python module named __init__.py file, return its directory. Otherwise return target unchanged.
+   * @param target PSI element to check
+   * @return PsiDirectory or target unchanged
+   */
   @Contract("null -> null; !null -> !null")
   @Nullable
   public static PsiElement turnInitIntoDir(@Nullable PsiElement target) {
@@ -977,7 +1003,7 @@ public class PyUtil {
   }
 
   public static boolean isPackage(@NotNull PsiDirectory directory, @Nullable PsiElement anchor) {
-    if (turnDirIntoInit(directory) != null) {
+    if (directory.findFile(PyNames.INIT_DOT_PY) != null) {
       return true;
     }
     final LanguageLevel level = anchor != null ?

python/src/com/jetbrains/python/refactoring/classes/PyClassRefactoringUtil.java

@@ -245,7 +245,7 @@ public final class PyClassRefactoringUtil {
     final String asName = node.getCopyableUserData(ENCODED_IMPORT_AS);
     final Boolean useFromImport = node.getCopyableUserData(ENCODED_USE_FROM_IMPORT);
     if (target instanceof PsiDirectory) {
-      target = (PsiNamedElement)PyUtil.turnDirIntoInit(target);
+      target = (PsiNamedElement)PyUtil.turnDirIntoPackageElement((PsiDirectory)target, node);
     }
     if (target instanceof PyFunction) {
       final PyFunction f = (PyFunction)target;
@@ -256,7 +256,7 @@ public final class PyClassRefactoringUtil {
     }
     if (target == null) return;
     if (PsiTreeUtil.isAncestor(node.getContainingFile(), target, false)) return;
-    if (target instanceof PyFile) {
+    if (target instanceof PyFile || target instanceof PsiDirectory) {
       insertImport(node, target, asName, useFromImport != null ? useFromImport : true);
     }
     else {
@@ -302,16 +302,16 @@ public final class PyClassRefactoringUtil {
                                      @Nullable String asName,
                                      boolean preferFromImport) {
     if (PyBuiltinCache.getInstance(element).isBuiltin(element)) return false;
-    final PsiFile newFile = element.getContainingFile();
+    final PsiFileSystemItem elementSource = element instanceof PsiDirectory? (PsiFileSystemItem)element : element.getContainingFile();
     final PsiFile file = anchor.getContainingFile();
-    if (newFile == file) return false;
+    if (elementSource == file) return false;
     final QualifiedName qname = QualifiedNameFinder.findCanonicalImportPath(element, anchor);
     if (qname == null || !isValidQualifiedName(qname)) {
       return false;
     }
     final QualifiedName containingQName;
     final String importedName;
-    if (element instanceof PyFile) {
+    if (element instanceof PyFile || element instanceof PsiDirectory) {
       containingQName = qname.removeLastComponent();
       importedName = qname.getLastComponent();
     }
@@ -319,7 +319,7 @@ public final class PyClassRefactoringUtil {
       containingQName = qname;
       importedName = getOriginalName(element);
     }
-    final AddImportHelper.ImportPriority priority = AddImportHelper.getImportPriority(anchor, newFile);
+    final AddImportHelper.ImportPriority priority = AddImportHelper.getImportPriority(anchor, elementSource);
     if (preferFromImport && !containingQName.getComponents().isEmpty()) {
       return AddImportHelper.addOrUpdateFromImportStatement(file, containingQName.toString(), importedName, asName, priority, anchor);
     }

python/testData/refactoring/move/namespacePackageUsedInMovedFunction/after/src/a.py

@@ -0,0 +1,2 @@
+
+

python/testData/refactoring/move/namespacePackageUsedInMovedFunction/after/src/b.py

@@ -0,0 +1,5 @@
+import nspkg
+
+
+def func():
+    print(nspkg)

python/testData/refactoring/move/namespacePackageUsedInMovedFunction/before/src/a.py

@@ -0,0 +1,4 @@
+import nspkg
+
+def func():
+    print(nspkg)

python/testData/refactoring/move/namespacePackageUsedInMovedFunction/before/src/b.py

@@ -0,0 +1,3 @@
+
+
+

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

@@ -199,6 +199,16 @@ public class PyMoveTest extends PyTestCase {
     doMoveFileTest("pkg1/subpkg1", "");
   }
 
+  // PY-14595
+  public void testNamespacePackageUsedInMovedFunction() {
+    runWithLanguageLevel(LanguageLevel.PYTHON33, new Runnable() {
+      @Override
+      public void run() {
+        doMoveSymbolTest("func", "b.py");
+      }
+    });
+  }
+
   public void testRelativeImportOfNameFromInitPy() {
     doMoveFileTest("pkg/subpkg2", "");
   }