From 2b64b8d52ed1c250310dc2e26cc22863cd9ad388 Mon Sep 17 00:00:00 2001
From: Dmitry Jemerov <yole@jetbrains.com>
Date: Wed, 25 Apr 2012 19:09:02 +0200
Subject: [PATCH] complete imported names inside __all__ (PY-6306)

---
 .../python/codeInsight/PyDunderAllReference.java   | 14 ++++++++++++--
 .../completion/dunderAllReferenceImport.after.py   |  3 +++
 .../completion/dunderAllReferenceImport.py         |  3 +++
 .../com/jetbrains/python/PythonCompletionTest.java |  4 ++++
 4 files changed, 22 insertions(+), 2 deletions(-)
 create mode 100644 python/testData/completion/dunderAllReferenceImport.after.py
 create mode 100644 python/testData/completion/dunderAllReferenceImport.py

diff --git a/python/src/com/jetbrains/python/codeInsight/PyDunderAllReference.java b/python/src/com/jetbrains/python/codeInsight/PyDunderAllReference.java
index b06a549f2f69..73e987c90e9e 100644
--- a/python/src/com/jetbrains/python/codeInsight/PyDunderAllReference.java
+++ b/python/src/com/jetbrains/python/codeInsight/PyDunderAllReference.java
@@ -1,5 +1,8 @@
 package com.jetbrains.python.codeInsight;
 
+import com.intellij.codeInsight.lookup.LookupElement;
+import com.intellij.codeInsight.lookup.LookupElementBuilder;
+import com.intellij.openapi.util.Iconable;
 import com.intellij.openapi.util.Key;
 import com.intellij.openapi.util.TextRange;
 import com.intellij.psi.PsiElement;
@@ -9,6 +12,7 @@ import com.intellij.psi.ResolveState;
 import com.intellij.psi.scope.PsiScopeProcessor;
 import com.intellij.util.ArrayUtil;
 import com.jetbrains.python.psi.PyFile;
+import com.jetbrains.python.psi.PyImportElement;
 import com.jetbrains.python.psi.PyStringLiteralExpression;
 import com.jetbrains.python.psi.PyUtil;
 import com.jetbrains.python.psi.impl.LightNamedElement;
@@ -41,7 +45,7 @@ public class PyDunderAllReference extends PsiReferenceBase<PyStringLiteralExpres
   @NotNull
   @Override
   public Object[] getVariants() {
-    final List<PsiElement> result = new ArrayList<PsiElement>();
+    final List<LookupElement> result = new ArrayList<LookupElement>();
     PyFile containingFile = (PyFile) getElement().getContainingFile().getOriginalFile();
     final List<String> dunderAll = containingFile.getDunderAll();
     containingFile.processDeclarations(new PsiScopeProcessor() {
@@ -50,7 +54,13 @@ public class PyDunderAllReference extends PsiReferenceBase<PyStringLiteralExpres
         if (element instanceof PsiNamedElement && !(element instanceof LightNamedElement)) {
           final String name = ((PsiNamedElement)element).getName();
           if (name != null && PyUtil.getInitialUnderscores(name) == 0 && (dunderAll == null || !dunderAll.contains(name))) {
-            result.add(element);
+            result.add(LookupElementBuilder.create((PsiNamedElement) element).setIcon(element.getIcon(Iconable.ICON_FLAG_CLOSED)));
+          }
+        }
+        else if (element instanceof PyImportElement) {
+          final String visibleName = ((PyImportElement)element).getVisibleName();
+          if (visibleName != null && (dunderAll == null || !dunderAll.contains(visibleName))) {
+            result.add(LookupElementBuilder.create(element, visibleName));
           }
         }
         return true;
diff --git a/python/testData/completion/dunderAllReferenceImport.after.py b/python/testData/completion/dunderAllReferenceImport.after.py
new file mode 100644
index 000000000000..2011874196d6
--- /dev/null
+++ b/python/testData/completion/dunderAllReferenceImport.after.py
@@ -0,0 +1,3 @@
+from sys import hexversion
+
+__all__ = ['hexversion']
diff --git a/python/testData/completion/dunderAllReferenceImport.py b/python/testData/completion/dunderAllReferenceImport.py
new file mode 100644
index 000000000000..d091b755ca27
--- /dev/null
+++ b/python/testData/completion/dunderAllReferenceImport.py
@@ -0,0 +1,3 @@
+from sys import hexversion
+
+__all__ = ['hexve<caret>']
diff --git a/python/testSrc/com/jetbrains/python/PythonCompletionTest.java b/python/testSrc/com/jetbrains/python/PythonCompletionTest.java
index 92a4a2599df7..1d8ba03e122a 100644
--- a/python/testSrc/com/jetbrains/python/PythonCompletionTest.java
+++ b/python/testSrc/com/jetbrains/python/PythonCompletionTest.java
@@ -529,6 +529,10 @@ public class PythonCompletionTest extends PyTestCase {
     doTest();
   }
 
+  public void testDunderAllReferenceImport() {  // PY-6306
+    doTest();
+  }
+
   public void testOldStyleClassAttributes() {
     doTest();
   }