From 6e2d427be85424b74c18d46e6920104f2be3147d Mon Sep 17 00:00:00 2001
From: Ekaterina Tuzova <Ekaterina.Tuzova@jetbrains.com>
Date: Tue, 8 Sep 2015 13:28:42 +0300
Subject: [PATCH] added resolve in module for PyCharm EDU

---
 .../resources/META-INF/plugin.xml             |  1 +
 .../edu/learning/PyStudyImportResolver.java   | 30 +++++++++++++++++++
 2 files changed, 31 insertions(+)
 create mode 100644 python/edu/interactive-learning-python/src/com/jetbrains/edu/learning/PyStudyImportResolver.java

diff --git a/python/edu/interactive-learning-python/resources/META-INF/plugin.xml b/python/edu/interactive-learning-python/resources/META-INF/plugin.xml
index 3fe9801d8ea4..34f5cf04a607 100644
--- a/python/edu/interactive-learning-python/resources/META-INF/plugin.xml
+++ b/python/edu/interactive-learning-python/resources/META-INF/plugin.xml
@@ -24,6 +24,7 @@
   <extensions defaultExtensionNs="Pythonid">
     <inspectionExtension implementation="com.jetbrains.edu.learning.highlighting.PyStudyInspectionExtension"/>
     <pyReferenceResolveProvider implementation="com.jetbrains.edu.learning.PyStudyReferenceResolveProvider"/>
+    <importResolver implementation="com.jetbrains.edu.learning.PyStudyImportResolver"/>
   </extensions>
   <extensions defaultExtensionNs="Edu">
     <StudyExecutor implementationClass="com.jetbrains.edu.learning.PyStudyExecutor" language="Python"/>
diff --git a/python/edu/interactive-learning-python/src/com/jetbrains/edu/learning/PyStudyImportResolver.java b/python/edu/interactive-learning-python/src/com/jetbrains/edu/learning/PyStudyImportResolver.java
new file mode 100644
index 000000000000..49916e779c97
--- /dev/null
+++ b/python/edu/interactive-learning-python/src/com/jetbrains/edu/learning/PyStudyImportResolver.java
@@ -0,0 +1,30 @@
+package com.jetbrains.edu.learning;
+
+import com.intellij.psi.PsiDirectory;
+import com.intellij.psi.PsiElement;
+import com.intellij.psi.PsiFile;
+import com.intellij.psi.util.QualifiedName;
+import com.jetbrains.python.psi.impl.PyImportResolver;
+import com.jetbrains.python.psi.resolve.QualifiedNameResolveContext;
+import org.jetbrains.annotations.Nullable;
+
+public class PyStudyImportResolver implements PyImportResolver {
+  @Nullable
+  public PsiElement resolveImportReference(QualifiedName name, QualifiedNameResolveContext context, boolean withRoots) {
+    if (StudyTaskManager.getInstance(context.getProject()).getCourse() == null) {
+      return null;
+    }
+    final String nameString = name.toString();
+    PsiFile containingFile = context.getFootholdFile();
+    if (containingFile == null) return null;
+
+    final PsiDirectory directory = containingFile.getContainingDirectory();
+    if (directory == null) return null;
+    final PsiFile file = directory.findFile(nameString + ".py");
+    if (file != null) {
+      return file;
+    }
+
+    return null;
+  }
+}