From 01f2002956c066a34b9f2dc01c41e8451c88963d Mon Sep 17 00:00:00 2001
From: "andrey.matveev" <andrey.matveev@jetbrains.com>
Date: Tue, 22 Dec 2020 12:41:45 +0700
Subject: [PATCH] PY-46054 Fix missing same directory completion variants

(cherry picked from commit fb17e9285a39846b93708fb33c9822bff0a8da35)

IJ-MR-4823

GitOrigin-RevId: 1e0b31e14ff2c1ee95358b7bf68576f413fed71c
---
 .../python/psi/resolve/ResolveImportUtil.java          |  4 +++-
 .../fromImportFromSameDirectoryModule/foo_bar/lib.py   |  2 ++
 .../foo_bar/sample.after.py                            |  1 +
 .../foo_bar/sample.py                                  |  1 +
 .../com/jetbrains/python/Py3CompletionTest.java        | 10 ++++++++++
 5 files changed, 17 insertions(+), 1 deletion(-)
 create mode 100644 python/testData/completion/fromImportFromSameDirectoryModule/foo_bar/lib.py
 create mode 100644 python/testData/completion/fromImportFromSameDirectoryModule/foo_bar/sample.after.py
 create mode 100644 python/testData/completion/fromImportFromSameDirectoryModule/foo_bar/sample.py

diff --git a/python/python-psi-impl/src/com/jetbrains/python/psi/resolve/ResolveImportUtil.java b/python/python-psi-impl/src/com/jetbrains/python/psi/resolve/ResolveImportUtil.java
index 061014fb8f00..4048e40ebb94 100644
--- a/python/python-psi-impl/src/com/jetbrains/python/psi/resolve/ResolveImportUtil.java
+++ b/python/python-psi-impl/src/com/jetbrains/python/psi/resolve/ResolveImportUtil.java
@@ -2,6 +2,7 @@
 package com.jetbrains.python.psi.resolve;
 
 import com.google.common.collect.Lists;
+import com.intellij.codeInsight.completion.CompletionUtilCoreImpl;
 import com.intellij.openapi.module.Module;
 import com.intellij.openapi.module.ModuleUtilCore;
 import com.intellij.openapi.projectRoots.Sdk;
@@ -47,7 +48,8 @@ public final class ResolveImportUtil {
       if (file instanceof PyFile) {
         final PyFile pyFile = (PyFile)file;
         if (pyFile.getLanguageLevel().isPy3K()) {
-          if (foothold.getManager().isInProject(foothold) && Registry.is("python.explicit.namespace.packages")) {
+          PsiElement originalFoothold = CompletionUtilCoreImpl.getOriginalOrSelf(foothold);
+          if (foothold.getManager().isInProject(originalFoothold) && Registry.is("python.explicit.namespace.packages")) {
             return false;
           }
           return true;
diff --git a/python/testData/completion/fromImportFromSameDirectoryModule/foo_bar/lib.py b/python/testData/completion/fromImportFromSameDirectoryModule/foo_bar/lib.py
new file mode 100644
index 000000000000..6584985db844
--- /dev/null
+++ b/python/testData/completion/fromImportFromSameDirectoryModule/foo_bar/lib.py
@@ -0,0 +1,2 @@
+def foo():
+    pass
\ No newline at end of file
diff --git a/python/testData/completion/fromImportFromSameDirectoryModule/foo_bar/sample.after.py b/python/testData/completion/fromImportFromSameDirectoryModule/foo_bar/sample.after.py
new file mode 100644
index 000000000000..fad688cff38c
--- /dev/null
+++ b/python/testData/completion/fromImportFromSameDirectoryModule/foo_bar/sample.after.py
@@ -0,0 +1 @@
+from lib import foo
\ No newline at end of file
diff --git a/python/testData/completion/fromImportFromSameDirectoryModule/foo_bar/sample.py b/python/testData/completion/fromImportFromSameDirectoryModule/foo_bar/sample.py
new file mode 100644
index 000000000000..90b1c3c4e201
--- /dev/null
+++ b/python/testData/completion/fromImportFromSameDirectoryModule/foo_bar/sample.py
@@ -0,0 +1 @@
+from lib import f<caret>
\ No newline at end of file
diff --git a/python/testSrc/com/jetbrains/python/Py3CompletionTest.java b/python/testSrc/com/jetbrains/python/Py3CompletionTest.java
index 97f101332635..f044730c75c0 100644
--- a/python/testSrc/com/jetbrains/python/Py3CompletionTest.java
+++ b/python/testSrc/com/jetbrains/python/Py3CompletionTest.java
@@ -491,6 +491,16 @@ public class Py3CompletionTest extends PyTestCase {
     });
   }
 
+  // PY-46054
+  public void testFromImportFromSameDirectoryModule() {
+    runWithLanguageLevel(LanguageLevel.getLatest(), () -> {
+      myFixture.copyDirectoryToProject(getTestName(true), "");
+      myFixture.configureByFile("foo_bar/sample.py");
+      myFixture.completeBasic();
+      myFixture.checkResultByFile(getTestName(true) + "/foo_bar/sample.after.py");
+    });
+  }
+
   @Override
   protected String getTestDataPath() {
     return super.getTestDataPath() + "/completion";