completion of instance attributes shows variants from __slots__ (PY-1211)

python/src/com/jetbrains/python/psi/impl/PyClassImpl.java

@@ -721,13 +721,18 @@ public class PyClassImpl extends PyPresentableElementImpl<PyClassStub> implement
   }
 
   public void processDeclarations(@NotNull PsiScopeProcessor processor) {
-    // class level
+    if (!processClassLevelDeclarations(processor)) return;
+    if (!processInstanceLevelDeclarations(processor)) return;
+    processor.execute(this, ResolveState.initial());
+  }
+
+  public boolean processClassLevelDeclarations(PsiScopeProcessor processor) {
     final PyClassStub stub = getStub();
     if (stub != null) {
       final List<StubElement> children = stub.getChildrenStubs();
       for (StubElement child : children) {
         if (!processor.execute(child.getPsi(), ResolveState.initial())) {
-          return;
+          return false;
         }
       }
     }
@@ -735,16 +740,14 @@ public class PyClassImpl extends PyPresentableElementImpl<PyClassStub> implement
       final PsiElement the_psi = getNode().getPsi();
       PyResolveUtil.treeCrawlUp(processor, true, the_psi, the_psi);
     }
+    return true;
+  }
 
-    // instance level
+  public boolean processInstanceLevelDeclarations(PsiScopeProcessor processor) {
     for(PyTargetExpression expr: getInstanceAttributes()) {
-      if (!processor.execute(expr, ResolveState.initial())) return;
+      if (!processor.execute(expr, ResolveState.initial())) return false;
     }
-    //
-    if (processor instanceof VariantsProcessor) {
-      return;
-    }
-    processor.execute(this, ResolveState.initial());
+    return true;
   }
 
   public int getTextOffset() {

python/src/com/jetbrains/python/psi/types/PyClassType.java

@@ -6,7 +6,6 @@ import com.intellij.openapi.extensions.Extensions;
 import com.intellij.openapi.project.Project;
 import com.intellij.openapi.util.Condition;
 import com.intellij.psi.PsiElement;
-import com.intellij.psi.ResolveState;
 import com.intellij.util.ProcessingContext;
 import com.intellij.util.SmartList;
 import com.jetbrains.python.codeInsight.PyDynamicMember;
@@ -145,9 +144,10 @@ public class PyClassType implements PyType {
   }
 
   public Object[] getCompletionVariants(final PyQualifiedExpression referenceExpression, ProcessingContext context) {
-    List<? extends PsiElement> classList = new ParentMatcher(PyClass.class).search(referenceExpression);
-    boolean withinOurClass = classList != null && classList.get(0) == this;
     Set<String> namesAlready = context.get(CTX_NAMES);
+    if (namesAlready == null) {
+      namesAlready = new HashSet<String>();
+    }
     List<Object> ret = new ArrayList<Object>();
     Condition<String> underscoreFilter = new PyUtil.UnderscoreFilter(PyUtil.getInitialUnderscores(referenceExpression.getName()));
     // from providers
@@ -159,23 +159,48 @@ public class PyClassType implements PyType {
         }
       }
     }
-    // from our own class
+
+    addOwnClassMembers(referenceExpression, namesAlready, ret, underscoreFilter);
+
+    addInheritedMembers(referenceExpression, context, ret);
+    
+    return ret.toArray();
+  }
+
+  private void addOwnClassMembers(PyQualifiedExpression referenceExpression,
+                                  Set<String> namesAlready,
+                                  List<Object> ret, Condition<String> underscoreFilter) {
+    List<? extends PsiElement> classList = new ParentMatcher(PyClass.class).search(referenceExpression);
+    boolean withinOurClass = classList != null && classList.get(0) == this;
+
     final VariantsProcessor processor = new VariantsProcessor(
       referenceExpression, new PyResolveUtil.FilterNotInstance(myClass), underscoreFilter
     );
-    ((PyClassImpl) myClass).processDeclarations(processor);
-    if (namesAlready != null) {
-      for (LookupElement le : processor.getResultList()) {
-        String name = le.getLookupString();
-        if (namesAlready.contains(name)) continue;
-        if (!withinOurClass && isClassPrivate(name)) continue;
-        namesAlready.add(name);
-        ret.add(le);
+    ((PyClassImpl) myClass).processClassLevelDeclarations(processor);
+
+    List<String> slots = myClass.isNewStyleClass() ? myClass.getSlots() : null;
+    if (slots != null) {
+      processor.setAllowedNames(slots);
+    }
+    ((PyClassImpl) myClass).processInstanceLevelDeclarations(processor);
+
+    for (LookupElement le : processor.getResultList()) {
+      String name = le.getLookupString();
+      if (namesAlready.contains(name)) continue;
+      if (!withinOurClass && isClassPrivate(name)) continue;
+      namesAlready.add(name);
+      ret.add(le);
+    }
+    if (slots != null) {
+      for (String name : slots) {
+        if (!namesAlready.contains(name)) {
+          ret.add(LookupElementBuilder.create(name));
+        }
       }
     }
-    else {
-      ret.addAll(processor.getResultList());
-    }
+  }
+
+  private void addInheritedMembers(PyQualifiedExpression referenceExpression, ProcessingContext context, List<Object> ret) {
     for (PyClass ancestor : myClass.getSuperClasses()) {
       Object[] ancestry = (new PyClassType(ancestor, true)).getCompletionVariants(referenceExpression, context);
       for (Object ob : ancestry) {
@@ -189,7 +214,6 @@ public class PyClassType implements PyType {
       }
       ret.addAll(Arrays.asList(ancestry));
     }
-    return ret.toArray();
   }
 
   private static boolean isClassPrivate(String lookup_string) {

python/testData/completion/slots.after.py

@@ -0,0 +1,5 @@
+class A(object):
+    __slots__ = ['foo', 'bar']
+
+a = A()
+a.foo

python/testData/completion/slots.py

@@ -0,0 +1,5 @@
+class A(object):
+    __slots__ = ['foo', 'bar']
+
+a = A()
+a.f<caret>

python/testSrc/com/jetbrains/python/PythonCompletionTest.java

@@ -160,4 +160,8 @@ public class PythonCompletionTest extends PyLightFixtureTestCase {
     myFixture.completeBasic();
     assertSameElements(myFixture.getLookupElementStrings(), Arrays.asList("my_foo", "my_bar"));
   }
+
+  public void testSlots() throws Exception {  // PY-1211
+    doTest();
+  }
 }