account for comments and line breaks in __init__ files declaring namespace package (PY-38434)

GitOrigin-RevId: 4a5a890113a8d8aa55053f6d96d946958476ea17

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

@@ -408,25 +408,28 @@ private fun isNamespacePackage(element: PsiElement): Boolean {
   if (element is PsiDirectory) {
     val level = PyUtil.getLanguageLevelForVirtualFile(element.project, element.virtualFile)
     val initFile = PyUtil.turnDirIntoInit(element) ?: return !level.isPython2
-    val initText = initFile.text ?: return false
-    return isNamespaceDeclaration(initText, pkgutilsInitPatterns) || isNamespaceDeclaration(initText, pkgResourcesInitPatterns)
+    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
+    }
   }
   return false
 }
 
-private val pkgResourcesInitPatterns = listOf(
-  Pattern.compile("^__import__\\(['\"]pkg_resources['\"]\\).declare_namespace\\(__name__\\)\\s*\$"),
-  Pattern.compile("^from pkg_resources import declare_namespace\\Rdeclare_namespace\\(__name__\\)\\s*\$"),
-  Pattern.compile("^import pkg_resources\\Rpkg_resources.declare_namespace\\(__name__\\)\\s*\$"))
+private val oneLineNamespaceDeclarations = listOf(
+  Pattern.compile("^__path__[ ]?=[ ]?__import__\\(['\"]pkgutil['\"]\\).extend_path\\(__path__, __name__\\).*"),
+  Pattern.compile("^__import__\\(['\"]pkg_resources['\"]\\).declare_namespace\\(__name__\\).*"))
 
-
-private val pkgutilsInitPatterns = listOf(
-  Pattern.compile("^__path__[ ]?=[ ]?__import__\\(['\"]pkgutil['\"]\\).extend_path\\(__path__, __name__\\)\\s*\$"),
-  Pattern.compile("^from pkgutil import extend_path\\R__path__[ ]?=[ ]?extend_path\\(__path__,[ ]?__name__\\)\\s*\$"),
-  Pattern.compile("^import pkgutil\\R__path__[ ]?=[ ]?pkgutil\\.extend_path\\(__path__,[ ]?__name__\\)\\s*\$"))
-
-private fun isNamespaceDeclaration(text: String, initTextPatterns: List<Pattern>): Boolean =
-  initTextPatterns.asSequence().map { it.matcher(text) }.any { it.matches() }
+private val multilineNamespaceDeclarations = listOf(
+  listOf(Pattern.compile("^from pkgutil import extend_path.*"), Pattern.compile("^__path__[ ]?=[ ]?extend_path\\(__path__,[ ]?__name__\\).*")),
+  listOf(Pattern.compile("^import pkgutil.*"), Pattern.compile("^__path__[ ]?=[ ]?pkgutil\\.extend_path\\(__path__,[ ]?__name__\\).*")),
+  listOf(Pattern.compile("^from pkg_resources import declare_namespace.*"), Pattern.compile("^declare_namespace\\(__name__\\).*")),
+  listOf(Pattern.compile("^import pkg_resources.*"), Pattern.compile("^pkg_resources.declare_namespace\\(__name__\\).*")))
 
 
 private fun isUserFile(element: PsiElement, module: Module?) =

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

@@ -1,2 +1,2 @@
-from pkgutil import extend_path
-__path__ = extend_path(__path__, __name__)
+import pkgutil
+__path__ = pkgutil.extend_path(__path__, __name__)

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

@@ -0,0 +1,5 @@
+# This is a comment in __init__ file
+# to check that we still understand it's a part of namespace package
+#
+
+__path__= __import__('pkgutil').extend_path(__path__, __name__) # can be on the same line

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

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

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

@@ -0,0 +1,7 @@
+# This is a comment in __init__ file
+
+from pkgutil import extend_path # can be on the same line
+
+__path__ = extend_path(__path__, __name__)
+
+# to check that we still understand it's a part of namespace package

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

@@ -478,23 +478,21 @@ public class PyMultiFileResolveTest extends PyMultiFileResolveTestCase {
 
   // PY-16688
   public void testPkgResourcesNamespace() {
-    myFixture.copyDirectoryToProject("pkgResourcesNamespace", "");
-    runWithSourceRoots(Lists.newArrayList(myFixture.findFileInTempDir("root1"), myFixture.findFileInTempDir("root2")), () -> {
-      final PsiFile psiFile = myFixture.configureByFile("root1/pkg/a.py");
-      final PsiReference ref = PyResolveTestCase.findReferenceByMarker(psiFile);
-      assertInstanceOf(ref, PsiPolyVariantReference.class);
-      final List<PsiElement> elements = PyUtil.multiResolveTopPriority((PsiPolyVariantReference)ref);
-      assertEquals(1, elements.size());
-      PsiFile root1 = myFixture.getPsiManager().findFile(myFixture.findFileInTempDir("root1/pkg/__init__.py"));
-      PsiFile root2 = myFixture.getPsiManager().findFile(myFixture.findFileInTempDir("root2/pkg/__init__.py"));
-      assertNotParsed(root1);
-      assertNotParsed(root2);
-    });
+    doTestResolveInNamespacePackage(getTestName(true));
   }
 
   // PY-23087
   public void testPkgutilNamespace() {
-    myFixture.copyDirectoryToProject("pkgutilNamespace", "");
+    doTestResolveInNamespacePackage(getTestName(true));
+  }
+
+  // PY-38434
+  public void testPkgutilNamespaceWithComments() {
+    doTestResolveInNamespacePackage(getTestName(true));
+  }
+
+  private void doTestResolveInNamespacePackage(String namespace) {
+    myFixture.copyDirectoryToProject(namespace, "");
     runWithSourceRoots(Lists.newArrayList(myFixture.findFileInTempDir("root1"), myFixture.findFileInTempDir("root2")), () -> {
       final PsiFile psiFile = myFixture.configureByFile("root1/pkg/a.py");
       final PsiReference ref = PyResolveTestCase.findReferenceByMarker(psiFile);