Goto Symbol works for file-level variable declarations (PY-564)

python/src/META-INF/python-plugin-common.xml

@@ -54,6 +54,7 @@
     <stubIndex implementation="com.jetbrains.python.psi.stubs.PyClassNameIndexInsensitive"/>
     <stubIndex implementation="com.jetbrains.python.psi.stubs.PyFunctionNameIndex"/>
     <stubIndex implementation="com.jetbrains.python.psi.stubs.PySuperClassIndex"/>
+    <stubIndex implementation="com.jetbrains.python.psi.stubs.PyVariableNameIndex"/>
 
     <fileTypeIndentOptionsProvider implementation="com.jetbrains.python.formatter.PyIndentOptionsProvider"/>
 

python/src/com/jetbrains/python/PyGotoSymbolContributor.java

@@ -4,13 +4,16 @@ import com.intellij.navigation.ChooseByNameContributor;
 import com.intellij.navigation.NavigationItem;
 import com.intellij.openapi.project.Project;
 import com.intellij.psi.search.GlobalSearchScope;
-import com.intellij.psi.search.ProjectScope;
 import com.intellij.psi.stubs.StubIndex;
 import com.intellij.util.ArrayUtil;
 import com.jetbrains.python.psi.stubs.PyClassNameIndex;
 import com.jetbrains.python.psi.stubs.PyFunctionNameIndex;
+import com.jetbrains.python.psi.stubs.PyVariableNameIndex;
 
-import java.util.*;
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
 
 /**
  * @author yole
@@ -20,6 +23,7 @@ public class PyGotoSymbolContributor implements ChooseByNameContributor {
     Set<String> symbols = new HashSet<String>();
     symbols.addAll(PyClassNameIndex.allKeys(project));
     symbols.addAll(StubIndex.getInstance().getAllKeys(PyFunctionNameIndex.KEY, project));
+    symbols.addAll(StubIndex.getInstance().getAllKeys(PyVariableNameIndex.KEY, project));
     return ArrayUtil.toStringArray(symbols);
   }
 
@@ -31,6 +35,7 @@ public class PyGotoSymbolContributor implements ChooseByNameContributor {
     List<NavigationItem> symbols = new ArrayList<NavigationItem>();
     symbols.addAll(PyClassNameIndex.find(name, project, scope));
     symbols.addAll(PyFunctionNameIndex.find(name, project, scope));
+    symbols.addAll(PyVariableNameIndex.find(name, project, scope));
 
     return symbols.toArray(new NavigationItem[symbols.size()]);
   }

python/src/com/jetbrains/python/psi/impl/stubs/PyTargetExpressionElementType.java

@@ -13,7 +13,9 @@ import com.jetbrains.python.PyElementTypes;
 import com.jetbrains.python.psi.*;
 import com.jetbrains.python.psi.impl.PyQualifiedName;
 import com.jetbrains.python.psi.impl.PyTargetExpressionImpl;
+import com.jetbrains.python.psi.stubs.PyFileStub;
 import com.jetbrains.python.psi.stubs.PyTargetExpressionStub;
+import com.jetbrains.python.psi.stubs.PyVariableNameIndex;
 
 import java.io.IOException;
 
@@ -115,6 +117,12 @@ public class PyTargetExpressionElementType extends PyStubElementType<PyTargetExp
 
   @Override
   public void indexStub(PyTargetExpressionStub stub, IndexSink sink) {
+    if (stub.getParentStub() instanceof PyFileStub) {
+      String name = stub.getName();
+      if (name != null && PyUtil.getInitialUnderscores(name) == 0) {
+        sink.occurrence(PyVariableNameIndex.KEY, name);
+      }
+    }
     for (CustomTargetExpressionStubType stubType : getCustomStubTypes()) {
       stubType.indexStub(stub, sink);
     }

python/src/com/jetbrains/python/psi/stubs/PyVariableNameIndex.java

@@ -0,0 +1,25 @@
+package com.jetbrains.python.psi.stubs;
+
+import com.intellij.openapi.project.Project;
+import com.intellij.psi.search.GlobalSearchScope;
+import com.intellij.psi.stubs.StringStubIndexExtension;
+import com.intellij.psi.stubs.StubIndex;
+import com.intellij.psi.stubs.StubIndexKey;
+import com.jetbrains.python.psi.PyTargetExpression;
+
+import java.util.Collection;
+
+/**
+ * @author yole
+ */
+public class PyVariableNameIndex extends StringStubIndexExtension<PyTargetExpression> {
+  public static final StubIndexKey<String, PyTargetExpression> KEY = StubIndexKey.createIndexKey("Py.variable.shortName");
+
+  public StubIndexKey<String, PyTargetExpression> getKey() {
+    return KEY;
+  }
+
+  public static Collection<PyTargetExpression> find(String name, Project project, GlobalSearchScope scope) {
+    return StubIndex.getInstance().get(KEY, name, project, scope);
+  }
+}

python/testData/stubs/VariableIndex.py

@@ -0,0 +1,9 @@
+xyzzy = 1
+
+__all__ = ['xyzzy']
+
+def foo():
+    shazam = 2
+
+class C:
+    boohoo = 3

python/testSrc/com/jetbrains/python/PyStubsTest.java

@@ -9,6 +9,7 @@ import com.intellij.openapi.project.Project;
 import com.intellij.openapi.vfs.VirtualFile;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFile;
+import com.intellij.psi.search.GlobalSearchScope;
 import com.intellij.psi.stubs.StubElement;
 import com.intellij.psi.stubs.StubUpdatingIndex;
 import com.intellij.testFramework.TestDataPath;
@@ -19,8 +20,10 @@ import com.jetbrains.python.psi.impl.PyFileImpl;
 import com.jetbrains.python.psi.impl.PyQualifiedName;
 import com.jetbrains.python.psi.impl.PythonLanguageLevelPusher;
 import com.jetbrains.python.psi.stubs.PyClassStub;
+import com.jetbrains.python.psi.stubs.PyVariableNameIndex;
 import com.jetbrains.python.toolbox.Maybe;
 
+import java.util.Collection;
 import java.util.List;
 
 @TestDataPath("$CONTENT_ROOT/../testData/stubs/")
@@ -269,6 +272,16 @@ public class PyStubsTest extends PyLightFixtureTestCase {
     assertNotParsed(file);
   }
 
+  public void testVariableIndex() {
+    getTestFile();
+    GlobalSearchScope scope = GlobalSearchScope.allScope(myFixture.getProject());
+    Collection<PyTargetExpression> result = PyVariableNameIndex.find("xyzzy", myFixture.getProject(), scope);
+    assertEquals(1, result.size());
+    assertEquals(0, PyVariableNameIndex.find("shazam", myFixture.getProject(), scope).size());
+    assertEquals(0, PyVariableNameIndex.find("boohoo", myFixture.getProject(), scope).size());
+    assertEquals(0, PyVariableNameIndex.find("__all__", myFixture.getProject(), scope).size());
+  }
+
   private PyFile getTestFile() {
     return getTestFile(getTestName(false) + ".py");
   }