From 58fc2e7020acb051a4753b26fc51f80c19e2ff3d Mon Sep 17 00:00:00 2001
From: Andrey Vlasovskikh <andrey.vlasovskikh@jetbrains.com>
Date: Mon, 13 May 2013 22:03:28 +0400
Subject: [PATCH] Use original file as origin in TypeEvalContext during
 completion

---
 .../python/documentation/DocStringTypeReference.java          | 3 ++-
 .../python/psi/impl/references/PyImportReference.java         | 3 ++-
 .../python/psi/impl/references/PyQualifiedReference.java      | 2 +-
 .../src/com/jetbrains/python/psi/types/PyClassTypeImpl.java   | 4 +++-
 4 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/python/src/com/jetbrains/python/documentation/DocStringTypeReference.java b/python/src/com/jetbrains/python/documentation/DocStringTypeReference.java
index 9b8dace47630..8a22ec626c41 100644
--- a/python/src/com/jetbrains/python/documentation/DocStringTypeReference.java
+++ b/python/src/com/jetbrains/python/documentation/DocStringTypeReference.java
@@ -1,6 +1,7 @@
 package com.jetbrains.python.documentation;
 
 import com.google.common.collect.Lists;
+import com.intellij.codeInsight.completion.CompletionUtil;
 import com.intellij.openapi.util.TextRange;
 import com.intellij.psi.*;
 import com.intellij.util.IncorrectOperationException;
@@ -113,7 +114,7 @@ public class DocStringTypeReference extends PsiPolyVariantReferenceBase<PsiEleme
         for (PyImportElement element : elements) {
           final PyReferenceExpression referenceExpression = element.getImportReferenceExpression();
           if (referenceExpression == null) continue;
-          final PyType type = TypeEvalContext.userInitiated(file).getType(referenceExpression);
+          final PyType type = TypeEvalContext.userInitiated(CompletionUtil.getOriginalOrSelf(file)).getType(referenceExpression);
           if (type instanceof PyClassType) {
             variants.add(((PyClassType)type).getPyClass());
           }
diff --git a/python/src/com/jetbrains/python/psi/impl/references/PyImportReference.java b/python/src/com/jetbrains/python/psi/impl/references/PyImportReference.java
index 4d4e14c779e1..934edbbdbcac 100644
--- a/python/src/com/jetbrains/python/psi/impl/references/PyImportReference.java
+++ b/python/src/com/jetbrains/python/psi/impl/references/PyImportReference.java
@@ -1,5 +1,6 @@
 package com.jetbrains.python.psi.impl.references;
 
+import com.intellij.codeInsight.completion.CompletionUtil;
 import com.intellij.codeInsight.completion.InsertHandler;
 import com.intellij.codeInsight.completion.InsertionContext;
 import com.intellij.codeInsight.lookup.LookupElement;
@@ -81,7 +82,7 @@ public class PyImportReference extends PyReferenceImpl {
     }
 
     PyExpression qualifier = myElement.getQualifier();
-    final TypeEvalContext context = TypeEvalContext.userInitiated(myElement.getContainingFile());
+    final TypeEvalContext context = TypeEvalContext.userInitiated(CompletionUtil.getOriginalOrSelf(myElement).getContainingFile());
     if (qualifier != null) {
       // qualifier's type must be module, it should know how to complete
       PyType type = context.getType(qualifier);
diff --git a/python/src/com/jetbrains/python/psi/impl/references/PyQualifiedReference.java b/python/src/com/jetbrains/python/psi/impl/references/PyQualifiedReference.java
index f1d900f49319..7e177b15df71 100644
--- a/python/src/com/jetbrains/python/psi/impl/references/PyQualifiedReference.java
+++ b/python/src/com/jetbrains/python/psi/impl/references/PyQualifiedReference.java
@@ -247,7 +247,7 @@ public class PyQualifiedReference extends PyReferenceImpl {
     }
     final PyQualifiedExpression element = CompletionUtil.getOriginalOrSelf(myElement);
 
-    PyType qualifierType = TypeEvalContext.userInitiated(myElement.getContainingFile()).getType(qualifier);
+    PyType qualifierType = TypeEvalContext.userInitiated(element.getContainingFile()).getType(qualifier);
     ProcessingContext ctx = new ProcessingContext();
     final Set<String> namesAlready = new HashSet<String>();
     ctx.put(PyType.CTX_NAMES, namesAlready);
diff --git a/python/src/com/jetbrains/python/psi/types/PyClassTypeImpl.java b/python/src/com/jetbrains/python/psi/types/PyClassTypeImpl.java
index 0d6e61c479d5..555acaeb49f9 100644
--- a/python/src/com/jetbrains/python/psi/types/PyClassTypeImpl.java
+++ b/python/src/com/jetbrains/python/psi/types/PyClassTypeImpl.java
@@ -348,7 +348,9 @@ public class PyClassTypeImpl extends UserDataHolderBase implements PyClassType {
     boolean suppressParentheses = context.get(CTX_SUPPRESS_PARENTHESES) != null;
     addOwnClassMembers(location, namesAlready, suppressParentheses, ret);
 
-    final TypeEvalContext typeEvalContext = TypeEvalContext.userInitiated(location != null ? location.getContainingFile() : null);
+    final TypeEvalContext typeEvalContext = TypeEvalContext.userInitiated(location != null ?
+                                                                          CompletionUtil.getOriginalOrSelf(location).getContainingFile() :
+                                                                          null);
     addInheritedMembers(prefix, location, namesAlready, context, ret, typeEvalContext);
 
     // from providers