From 3490280116e18b06d81ddcc68836ec375faad50c Mon Sep 17 00:00:00 2001
From: Dmitry Jemerov <yole@jetbrains.com>
Date: Wed, 22 Aug 2012 15:42:48 +0200
Subject: [PATCH] a PyTargetExpression will itself call
 PyTypeProvider.getReferenceType() when asked for its type (PY-7270)

---
 .../controlflow/InstructionTypeCallback.java       |  1 +
 .../controlflow/ReadWriteInstruction.java          | 12 +++++++++++-
 .../python/psi/impl/PyReferenceExpressionImpl.java | 14 +++++++++-----
 .../python/psi/impl/PyTargetExpressionImpl.java    |  9 +++++++++
 4 files changed, 30 insertions(+), 6 deletions(-)

diff --git a/python/src/com/jetbrains/python/codeInsight/controlflow/InstructionTypeCallback.java b/python/src/com/jetbrains/python/codeInsight/controlflow/InstructionTypeCallback.java
index d65637dcb2d8..36c4cfd689fb 100644
--- a/python/src/com/jetbrains/python/codeInsight/controlflow/InstructionTypeCallback.java
+++ b/python/src/com/jetbrains/python/codeInsight/controlflow/InstructionTypeCallback.java
@@ -9,5 +9,6 @@ import org.jetbrains.annotations.Nullable;
  * @author yole
  */
 public interface InstructionTypeCallback {
+  @Nullable
   PyType getType(TypeEvalContext context, @Nullable PsiElement anchor);
 }
diff --git a/python/src/com/jetbrains/python/codeInsight/controlflow/ReadWriteInstruction.java b/python/src/com/jetbrains/python/codeInsight/controlflow/ReadWriteInstruction.java
index a936acb73425..57f1d64f3593 100644
--- a/python/src/com/jetbrains/python/codeInsight/controlflow/ReadWriteInstruction.java
+++ b/python/src/com/jetbrains/python/codeInsight/controlflow/ReadWriteInstruction.java
@@ -5,6 +5,7 @@ import com.intellij.codeInsight.controlflow.impl.InstructionImpl;
 import com.intellij.psi.PsiElement;
 import com.jetbrains.python.psi.PyElement;
 import com.jetbrains.python.psi.PyExpression;
+import com.jetbrains.python.psi.impl.PyTargetExpressionImpl;
 import com.jetbrains.python.psi.types.PyType;
 import com.jetbrains.python.psi.types.TypeEvalContext;
 import org.jetbrains.annotations.NonNls;
@@ -12,9 +13,18 @@ import org.jetbrains.annotations.Nullable;
 
 public class ReadWriteInstruction extends InstructionImpl {
   final InstructionTypeCallback EXPR_TYPE = new InstructionTypeCallback() {
+    @Nullable
     @Override
     public PyType getType(TypeEvalContext context, @Nullable PsiElement anchor) {
-      return myElement instanceof PyExpression ? context.getType((PyExpression)myElement) : null;
+      if (myElement instanceof PyTargetExpressionImpl) {
+        return ((PyTargetExpressionImpl) myElement).getTypeWithAnchor(context, anchor);
+      }
+      if (myElement instanceof PyExpression) {
+        return context.getType((PyExpression)myElement);
+      }
+      else {
+        return null;
+      }
     }
   };
 
diff --git a/python/src/com/jetbrains/python/psi/impl/PyReferenceExpressionImpl.java b/python/src/com/jetbrains/python/psi/impl/PyReferenceExpressionImpl.java
index 36a04c0471a9..e1b56721082c 100644
--- a/python/src/com/jetbrains/python/psi/impl/PyReferenceExpressionImpl.java
+++ b/python/src/com/jetbrains/python/psi/impl/PyReferenceExpressionImpl.java
@@ -281,9 +281,11 @@ public class PyReferenceExpressionImpl extends PyElementImpl implements PyRefere
   public static PyType getTypeFromTarget(@NotNull final PsiElement target,
                                          final TypeEvalContext context,
                                          PyReferenceExpression anchor) {
-    final PyType pyType = getReferenceTypeFromProviders(target, context, anchor);
-    if (pyType != null) {
-      return pyType;
+    if (!(target instanceof PyTargetExpression)) {  // PyTargetExpression will ask about its type itself
+      final PyType pyType = getReferenceTypeFromProviders(target, context, anchor);
+      if (pyType != null) {
+        return pyType;
+      }
     }
     if (target instanceof PyTargetExpression) {
       final String name = ((PyTargetExpression)target).getName();
@@ -329,7 +331,7 @@ public class PyReferenceExpressionImpl extends PyElementImpl implements PyRefere
         if (propertyDecorator != null) {
           return PyBuiltinCache.getInstance(target).getObjectType(PyNames.PROPERTY);
         }
-        for (PyDecorator decorator: decoratorList.getDecorators()) {
+        for (PyDecorator decorator : decoratorList.getDecorators()) {
           final PyQualifiedName qName = decorator.getQualifiedName();
           if (qName != null && (qName.endsWith(PyNames.SETTER) || qName.endsWith(PyNames.DELETER))) {
             return PyBuiltinCache.getInstance(target).getObjectType(PyNames.PROPERTY);
@@ -342,7 +344,9 @@ public class PyReferenceExpressionImpl extends PyElementImpl implements PyRefere
     }
     if (target instanceof PsiDirectory) {
       PsiFile file = ((PsiDirectory)target).findFile(PyNames.INIT_DOT_PY);
-      if (file != null) return getTypeFromTarget(file, context, anchor);
+      if (file != null) {
+        return getTypeFromTarget(file, context, anchor);
+      }
     }
     return null;
   }
diff --git a/python/src/com/jetbrains/python/psi/impl/PyTargetExpressionImpl.java b/python/src/com/jetbrains/python/psi/impl/PyTargetExpressionImpl.java
index cb39aea1dd47..eafcb179fd1f 100644
--- a/python/src/com/jetbrains/python/psi/impl/PyTargetExpressionImpl.java
+++ b/python/src/com/jetbrains/python/psi/impl/PyTargetExpressionImpl.java
@@ -98,6 +98,11 @@ public class PyTargetExpressionImpl extends PyPresentableElementImpl<PyTargetExp
   }
 
   public PyType getType(@NotNull TypeEvalContext context) {
+    return getTypeWithAnchor(context, null);
+  }
+
+  @Nullable
+  public PyType getTypeWithAnchor(TypeEvalContext context, @Nullable PsiElement anchor) {
     if (!TypeEvalStack.mayEvaluate(this)) {
       return null;
     }
@@ -107,6 +112,10 @@ public class PyTargetExpressionImpl extends PyPresentableElementImpl<PyTargetExp
         // imported via __all__
         return null;
       }
+      final PyType pyType = PyReferenceExpressionImpl.getReferenceTypeFromProviders(this, context, anchor);
+      if (pyType != null) {
+        return pyType;
+      }
       if (!context.maySwitchToAST(this)) {
         final PsiElement value = getStub() != null ? findAssignedValueByStub() : findAssignedValue();
         if (value instanceof PyTypedElement) {