diff --git a/python/src/com/jetbrains/python/psi/PyUtil.java b/python/src/com/jetbrains/python/psi/PyUtil.java index 5946be8705da..19d6a7e511fe 100644 --- a/python/src/com/jetbrains/python/psi/PyUtil.java +++ b/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 ? diff --git a/python/src/com/jetbrains/python/refactoring/classes/PyClassRefactoringUtil.java b/python/src/com/jetbrains/python/refactoring/classes/PyClassRefactoringUtil.java index 458fc127b29e..c748731559b0 100644 --- a/python/src/com/jetbrains/python/refactoring/classes/PyClassRefactoringUtil.java +++ b/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); } diff --git a/python/testData/refactoring/move/namespacePackageUsedInMovedFunction/after/src/a.py b/python/testData/refactoring/move/namespacePackageUsedInMovedFunction/after/src/a.py new file mode 100644 index 000000000000..139597f9cb07 --- /dev/null +++ b/python/testData/refactoring/move/namespacePackageUsedInMovedFunction/after/src/a.py @@ -0,0 +1,2 @@ + + diff --git a/python/testData/refactoring/move/namespacePackageUsedInMovedFunction/after/src/b.py b/python/testData/refactoring/move/namespacePackageUsedInMovedFunction/after/src/b.py new file mode 100644 index 000000000000..3ed2229dba37 --- /dev/null +++ b/python/testData/refactoring/move/namespacePackageUsedInMovedFunction/after/src/b.py @@ -0,0 +1,5 @@ +import nspkg + + +def func(): + print(nspkg) \ No newline at end of file diff --git a/python/testData/refactoring/move/namespacePackageUsedInMovedFunction/after/src/nspkg/mod.py b/python/testData/refactoring/move/namespacePackageUsedInMovedFunction/after/src/nspkg/mod.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/python/testData/refactoring/move/namespacePackageUsedInMovedFunction/before/src/a.py b/python/testData/refactoring/move/namespacePackageUsedInMovedFunction/before/src/a.py new file mode 100644 index 000000000000..3b8b96056e5f --- /dev/null +++ b/python/testData/refactoring/move/namespacePackageUsedInMovedFunction/before/src/a.py @@ -0,0 +1,4 @@ +import nspkg + +def func(): + print(nspkg) \ No newline at end of file diff --git a/python/testData/refactoring/move/namespacePackageUsedInMovedFunction/before/src/b.py b/python/testData/refactoring/move/namespacePackageUsedInMovedFunction/before/src/b.py new file mode 100644 index 000000000000..b28b04f64312 --- /dev/null +++ b/python/testData/refactoring/move/namespacePackageUsedInMovedFunction/before/src/b.py @@ -0,0 +1,3 @@ + + + diff --git a/python/testData/refactoring/move/namespacePackageUsedInMovedFunction/before/src/nspkg/mod.py b/python/testData/refactoring/move/namespacePackageUsedInMovedFunction/before/src/nspkg/mod.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/python/testSrc/com/jetbrains/python/refactoring/PyMoveTest.java b/python/testSrc/com/jetbrains/python/refactoring/PyMoveTest.java index b96d827c3342..fbe11b8572a8 100644 --- a/python/testSrc/com/jetbrains/python/refactoring/PyMoveTest.java +++ b/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", ""); }