From 36ebf9ef46f5e69f047ebb07ea6ec89f1b9e1f2e Mon Sep 17 00:00:00 2001
From: Dmitry Jemerov <yole@jetbrains.com>
Date: Thu, 30 Aug 2012 19:22:17 +0200
Subject: [PATCH] findReferencesToHighlight() correctly handles dir/__init__.py
 duality (PY-1514)

---
 .../findUsages/PyModuleFindUsagesHandler.java | 20 +++++++++++++++++++
 1 file changed, 20 insertions(+)

diff --git a/python/src/com/jetbrains/python/findUsages/PyModuleFindUsagesHandler.java b/python/src/com/jetbrains/python/findUsages/PyModuleFindUsagesHandler.java
index 9441948095bb..689c77aa644d 100644
--- a/python/src/com/jetbrains/python/findUsages/PyModuleFindUsagesHandler.java
+++ b/python/src/com/jetbrains/python/findUsages/PyModuleFindUsagesHandler.java
@@ -6,11 +6,20 @@ import com.intellij.find.findUsages.FindUsagesHandler;
 import com.intellij.psi.PsiDirectory;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFileSystemItem;
+import com.intellij.psi.PsiReference;
+import com.intellij.psi.search.SearchScope;
+import com.intellij.psi.search.searches.ReferencesSearch;
 import com.intellij.ui.SimpleColoredComponent;
 import com.intellij.ui.SimpleTextAttributes;
+import com.jetbrains.python.PyNames;
+import com.jetbrains.python.psi.PyFile;
 import com.jetbrains.python.psi.PyUtil;
 import org.jetbrains.annotations.NotNull;
 
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.List;
+
 /**
  * @author yole
  */
@@ -46,4 +55,15 @@ public class PyModuleFindUsagesHandler extends FindUsagesHandler {
       }
     };
   }
+
+  @Override
+  public Collection<PsiReference> findReferencesToHighlight(PsiElement target, SearchScope searchScope) {
+    if (target instanceof PyFile && PyNames.INIT_DOT_PY.equals(((PyFile)target).getName())) {
+      List<PsiReference> result = new ArrayList<PsiReference>();
+      result.addAll(super.findReferencesToHighlight(target, searchScope));
+      result.addAll(ReferencesSearch.search(PyUtil.turnInitIntoDir(target), searchScope, false).findAll());
+      return result;
+    }
+    return super.findReferencesToHighlight(target, searchScope);
+  }
 }