diff --git a/python/python-psi-impl/src/com/jetbrains/python/psi/resolve/PyResolveImportUtil.kt b/python/python-psi-impl/src/com/jetbrains/python/psi/resolve/PyResolveImportUtil.kt index b63dbf8a1c53..992b209b061a 100644 --- a/python/python-psi-impl/src/com/jetbrains/python/psi/resolve/PyResolveImportUtil.kt +++ b/python/python-psi-impl/src/com/jetbrains/python/psi/resolve/PyResolveImportUtil.kt @@ -405,22 +405,41 @@ private fun isNamespacePackage(element: PsiElement): Boolean { if (element is PsiDirectory) { val level = LanguageLevel.forElement(element) val initFile = PyUtil.turnDirIntoInit(element) ?: return !level.isPython2 - val initLines = initFile.text.lineSequence() - .filterNot { line -> line.trim().let { it.isEmpty() || it.startsWith("#") } } - .take(2) - .toList() - return when (initLines.size) { - 1 -> oneLineNamespaceDeclarations.any { it.matcher(initLines.first()).matches() } - 2 -> multilineNamespaceDeclarations.any { it[0].matcher(initLines[0]).matches() && it[1].matcher(initLines[1]).matches() } - else -> false + + var nextMultiLinePattern: Pattern? = null + var nextTryPattern: Pattern? = null + var inDocstring = false + var afterTry = false + + loop@ for (line in initFile.text.lineSequence()) { + when { + inDocstring -> inDocstring = !line.endsWith("\"\"\"") + line.isBlank() || line.trim().startsWith("#") || line.startsWith("except") -> continue@loop + line.startsWith("\"\"\"") -> inDocstring = true + line.startsWith("try:") -> afterTry = true + + afterTry && nextTryPattern != null -> return nextTryPattern.matcher(line).matches() + afterTry -> nextTryPattern = when { + oneLineNamespaceDeclarations[0].matcher(line).matches() -> oneLineNamespaceDeclarations[1] + oneLineNamespaceDeclarations[1].matcher(line).matches() -> oneLineNamespaceDeclarations[0] + else -> return false + } + + oneLineNamespaceDeclarations.any { it.matcher(line).matches()} -> return true + nextMultiLinePattern == null -> { + nextMultiLinePattern = multilineNamespaceDeclarations.find { it[0].matcher(line).matches() }?.get(1) ?: return false + } + nextMultiLinePattern.matcher(line).matches() -> return true + else -> return false + } } } return false } private val oneLineNamespaceDeclarations = listOf( - Pattern.compile("^__path__[ ]?=[ ]?__import__\\(['\"]pkgutil['\"]\\).extend_path\\(__path__, __name__\\).*"), - Pattern.compile("^__import__\\(['\"]pkg_resources['\"]\\).declare_namespace\\(__name__\\).*")) + Pattern.compile("\\s*__path__[ ]?=[ ]?__import__\\(['\"]pkgutil['\"]\\).extend_path\\(__path__, __name__\\).*"), + Pattern.compile("\\s*__import__\\(['\"]pkg_resources['\"]\\).declare_namespace\\(__name__\\).*")) private val multilineNamespaceDeclarations = listOf( listOf(Pattern.compile("^from pkgutil import extend_path.*"), Pattern.compile("^__path__[ ]?=[ ]?extend_path\\(__path__,[ ]?__name__\\).*")), diff --git a/python/testData/resolve/multiFile/pkgResourcesNamespaceWithDocstring/root1/pkg/__init__.py b/python/testData/resolve/multiFile/pkgResourcesNamespaceWithDocstring/root1/pkg/__init__.py new file mode 100644 index 000000000000..5619c361a50e --- /dev/null +++ b/python/testData/resolve/multiFile/pkgResourcesNamespaceWithDocstring/root1/pkg/__init__.py @@ -0,0 +1,5 @@ +""" +Docstring here should have no effect on resolve. +""" +import pkg_resources +pkg_resources.declare_namespace(__name__) diff --git a/python/testData/resolve/multiFile/pkgResourcesNamespaceWithDocstring/root1/pkg/a.py b/python/testData/resolve/multiFile/pkgResourcesNamespaceWithDocstring/root1/pkg/a.py new file mode 100644 index 000000000000..c394838ae54d --- /dev/null +++ b/python/testData/resolve/multiFile/pkgResourcesNamespaceWithDocstring/root1/pkg/a.py @@ -0,0 +1,2 @@ +import pkg.second +# diff --git a/python/testData/resolve/multiFile/pkgResourcesNamespaceWithDocstring/root2/pkg/__init__.py b/python/testData/resolve/multiFile/pkgResourcesNamespaceWithDocstring/root2/pkg/__init__.py new file mode 100644 index 000000000000..39ca03452282 --- /dev/null +++ b/python/testData/resolve/multiFile/pkgResourcesNamespaceWithDocstring/root2/pkg/__init__.py @@ -0,0 +1,4 @@ +""" +Docstring here should have no effect on resolve. +""" +__import__('pkg_resources').declare_namespace(__name__) diff --git a/python/testData/resolve/multiFile/pkgResourcesNamespaceWithDocstring/root2/pkg/second.py b/python/testData/resolve/multiFile/pkgResourcesNamespaceWithDocstring/root2/pkg/second.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/python/testData/resolve/multiFile/tryExceptNamespace/root1/pkg/__init__.py b/python/testData/resolve/multiFile/tryExceptNamespace/root1/pkg/__init__.py new file mode 100644 index 000000000000..2e33574e7a2d --- /dev/null +++ b/python/testData/resolve/multiFile/tryExceptNamespace/root1/pkg/__init__.py @@ -0,0 +1,4 @@ +try: + __import__('pkg_resources').declare_namespace(__name__) +except ImportError: + __path__ = __import__('pkgutil').extend_path(__path__, __name__) \ No newline at end of file diff --git a/python/testData/resolve/multiFile/tryExceptNamespace/root1/pkg/a.py b/python/testData/resolve/multiFile/tryExceptNamespace/root1/pkg/a.py new file mode 100644 index 000000000000..c394838ae54d --- /dev/null +++ b/python/testData/resolve/multiFile/tryExceptNamespace/root1/pkg/a.py @@ -0,0 +1,2 @@ +import pkg.second +# diff --git a/python/testData/resolve/multiFile/tryExceptNamespace/root2/pkg/__init__.py b/python/testData/resolve/multiFile/tryExceptNamespace/root2/pkg/__init__.py new file mode 100644 index 000000000000..fc5352044b3b --- /dev/null +++ b/python/testData/resolve/multiFile/tryExceptNamespace/root2/pkg/__init__.py @@ -0,0 +1,9 @@ +""" +This is a docstring, that should be skipped. +""" +# Comments as well +try: + # even here + __import__('pkg_resources').declare_namespace(__name__) +except ImportError: + __path__ = __import__('pkgutil').extend_path(__path__, __name__) \ No newline at end of file diff --git a/python/testData/resolve/multiFile/tryExceptNamespace/root2/pkg/second.py b/python/testData/resolve/multiFile/tryExceptNamespace/root2/pkg/second.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/python/testSrc/com/jetbrains/python/PyMultiFileResolveTest.java b/python/testSrc/com/jetbrains/python/PyMultiFileResolveTest.java index 932fa8573afc..93fb14373448 100644 --- a/python/testSrc/com/jetbrains/python/PyMultiFileResolveTest.java +++ b/python/testSrc/com/jetbrains/python/PyMultiFileResolveTest.java @@ -491,6 +491,16 @@ public class PyMultiFileResolveTest extends PyMultiFileResolveTestCase { doTestResolveInNamespacePackage(getTestName(true)); } + // PY-39748 + public void testPkgResourcesNamespaceWithDocstring() { + doTestResolveInNamespacePackage(getTestName(true)); + } + + // PY-39748 + public void testTryExceptNamespace() { + doTestResolveInNamespacePackage(getTestName(true)); + } + private void doTestResolveInNamespacePackage(String namespace) { myFixture.copyDirectoryToProject(namespace, ""); runWithSourceRoots(Lists.newArrayList(myFixture.findFileInTempDir("root1"), myFixture.findFileInTempDir("root2")), () -> {