IDEA-CR-64748: fixes for namespace package resolve (PY-39748)

(cherry picked from commit c600a4b23556442cfcf7734dc212731017bc86a2)

GitOrigin-RevId: ac9d9b8c0a86f07533acbc4bba047a475581677b

python/python-psi-impl/src/com/jetbrains/python/psi/resolve/PyNamespacePackageUtil.kt

@@ -3,6 +3,7 @@ package com.jetbrains.python.psi.resolve
 
 import com.intellij.psi.PsiDirectory
 import com.intellij.psi.PsiElement
+import com.intellij.psi.tree.TokenSet
 import com.jetbrains.python.PyTokenTypes
 import com.jetbrains.python.lexer.PythonLexer
 import com.jetbrains.python.psi.LanguageLevel
@@ -17,8 +18,7 @@ fun isNamespacePackage(element: PsiElement): Boolean {
 
     val lexer = PythonLexer()
     lexer.start(initFile.text)
-    lexer.tokenType
-    while (lexer.tokenType in tokensToSkip) {
+    while (lexer.tokenType in TOKENS_TO_SKIP) {
       lexer.advance()
     }
 
@@ -26,31 +26,33 @@ fun isNamespacePackage(element: PsiElement): Boolean {
     var nextPattern: Pattern? = null
     for (line in codeStart.lineSequence()) {
       val trimmed = line.trim()
-      if (trimmed.startsWith("#")) continue
-      else if (nextPattern != null && nextPattern.matcher(trimmed).matches()) return true
-      else if (oneLineNamespaceDeclarations.any { it.matcher(trimmed).matches() }) return true
-      else if (nextPattern == null) nextPattern = multilineNamespaceDeclarations.find { it[0].matcher(trimmed).matches() }?.get(1)
-      else return false
+      if (trimmed.isEmpty() || trimmed.startsWith("#")) continue
+      else if (nextPattern != null) return nextPattern.matcher(trimmed).matches()
+      else if (ONE_LINE_NAMESPACE_DECLARATIONS.any { it.matcher(trimmed).matches() }) return true
+      else {
+        val matched = TWO_LINE_NAMESPACE_DECLARATIONS.find { it[0].matcher(trimmed).matches() }
+        nextPattern = matched?.get(1) ?: return false
+      }
     }
   }
   return false
 }
 
-private val tokensToSkip = setOf(PyTokenTypes.DOCSTRING,
-                                 PyTokenTypes.END_OF_LINE_COMMENT,
-                                 PyTokenTypes.LINE_BREAK,
-                                 PyTokenTypes.SPACE,
-                                 PyTokenTypes.TRY_KEYWORD,
-                                 PyTokenTypes.COLON)
+private val TOKENS_TO_SKIP = TokenSet.create(PyTokenTypes.DOCSTRING,
+                                             PyTokenTypes.END_OF_LINE_COMMENT,
+                                             PyTokenTypes.LINE_BREAK,
+                                             PyTokenTypes.SPACE,
+                                             PyTokenTypes.TRY_KEYWORD,
+                                             PyTokenTypes.COLON)
 
-private val multilineNamespaceDeclarations = listOf(
+private val TWO_LINE_NAMESPACE_DECLARATIONS = listOf(
   patterns("^from pkgutil import extend_path.*", "^__path__[ ]?=[ ]?extend_path\\(__path__,[ ]?__name__\\).*"),
   patterns("^import pkgutil.*", "^__path__[ ]?=[ ]?pkgutil\\.extend_path\\(__path__,[ ]?__name__\\).*"),
   patterns("^from pkg_resources import declare_namespace.*", "^declare_namespace\\(__name__\\).*"),
   patterns("^import pkg_resources.*", "^pkg_resources.declare_namespace\\(__name__\\).*")
 )
 
-private val oneLineNamespaceDeclarations = patterns(
+private val ONE_LINE_NAMESPACE_DECLARATIONS = patterns(
   "^__path__[ ]?=[ ]?__import__\\(['\"]pkgutil['\"]\\).extend_path\\(__path__, __name__\\).*",
   "^__import__\\(['\"]pkg_resources['\"]\\).declare_namespace\\(__name__\\).*"
 )

python/testData/resolve/multiFile/tryExceptMultilineNamespace/root1/pkg/__init__.py

@@ -0,0 +1,5 @@
+try:
+    import pkg_resources
+    pkg_resources.declare_namespace(__name__)
+except ImportError:
+    __path__ = __import__('pkgutil').extend_path(__path__, __name__)

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

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

python/testData/resolve/multiFile/tryExceptMultilineNamespace/root2/pkg/__init__.py

@@ -0,0 +1,12 @@
+"""
+This is a docstring, that should be skipped.
+"""
+# Comments as well
+try:
+    # even here
+    from pkg_resources import declare_namespace
+    # and here
+
+    declare_namespace(__name__)
+except ImportError:
+    __path__ = __import__('pkgutil').extend_path(__path__, __name__)

python/testSrc/com/jetbrains/python/PyMultiFileResolveTest.java

@@ -501,6 +501,11 @@ public class PyMultiFileResolveTest extends PyMultiFileResolveTestCase {
     doTestResolveInNamespacePackage(getTestName(true));
   }
 
+  // PY-39748
+  public void testTryExceptMultilineNamespace() {
+    doTestResolveInNamespacePackage(getTestName(true));
+  }
+
   private void doTestResolveInNamespacePackage(String namespace) {
     myFixture.copyDirectoryToProject(namespace, "");
     runWithSourceRoots(Lists.newArrayList(myFixture.findFileInTempDir("root1"), myFixture.findFileInTempDir("root2")), () -> {