diff --git a/python/openapi/src/com/jetbrains/python/templateLanguages/TemplateFileReferenceSet.java b/python/openapi/src/com/jetbrains/python/templateLanguages/TemplateFileReferenceSet.java
index ff554151d344..0b0a719df2f6 100644
--- a/python/openapi/src/com/jetbrains/python/templateLanguages/TemplateFileReferenceSet.java
+++ b/python/openapi/src/com/jetbrains/python/templateLanguages/TemplateFileReferenceSet.java
@@ -6,7 +6,10 @@ import com.intellij.openapi.util.Pair;
 import com.intellij.openapi.util.SystemInfo;
 import com.intellij.openapi.util.TextRange;
 import com.intellij.openapi.vfs.VirtualFile;
-import com.intellij.psi.*;
+import com.intellij.psi.PsiElement;
+import com.intellij.psi.PsiFileSystemItem;
+import com.intellij.psi.PsiManager;
+import com.intellij.psi.PsiReferenceProvider;
 import com.intellij.psi.impl.source.resolve.reference.impl.providers.FileReference;
 import com.intellij.psi.impl.source.resolve.reference.impl.providers.FileReferenceSet;
 import com.intellij.util.containers.ContainerUtil;
@@ -82,7 +85,7 @@ public class TemplateFileReferenceSet extends FileReferenceSet {
     if (module != null) {
       List<VirtualFile> templatesFolders = getRoots(module);
       for (VirtualFile folder : templatesFolders) {
-        final PsiDirectory directory = PsiManager.getInstance(module.getProject()).findDirectory(folder);
+        final PsiFileSystemItem directory = getPsiDirectory(module, folder);
         if (directory != null) {
           contexts.add(directory);
         }
@@ -91,6 +94,10 @@ public class TemplateFileReferenceSet extends FileReferenceSet {
     return contexts;
   }
 
+  protected PsiFileSystemItem getPsiDirectory(Module module, VirtualFile folder) {
+    return PsiManager.getInstance(module.getProject()).findDirectory(folder);
+  }
+
   protected List<VirtualFile> getRoots(Module module) {
     return TemplatesService.getInstance(module).getTemplateFolders();
   }