IDEA-CR-64748: account for docstring when resolving namespace package, support try-except style declarations (PY-39748, PY-39512)

(cherry picked from commit 94907cff458e56607c11163bb681988a44750bc0)

GitOrigin-RevId: ff531a0b6f8fbf707656004d0e3ace38361e5cad

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__\\).*")),

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__)

python/testData/resolve/multiFile/pkgResourcesNamespaceWithDocstring/root1/pkg/a.py

@@ -0,0 +1,2 @@
+import pkg.second
+#          <ref>

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__)

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__)

python/testData/resolve/multiFile/tryExceptNamespace/root1/pkg/a.py

@@ -0,0 +1,2 @@
+import pkg.second
+#          <ref>

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__)

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")), () -> {