From b8d75dc401a9ddc03096c92f0c4a611a219e6d85 Mon Sep 17 00:00:00 2001
From: Andrey Vlasovskikh <andrey.vlasovskikh@gmail.com>
Date: Tue, 19 Aug 2014 00:14:18 +0400
Subject: [PATCH] Refactored PyAnnotation into a PyTypedElement

---
 .../jetbrains/python/psi/PyAnnotation.java    |  5 +---
 .../python/psi/impl/PyAnnotationImpl.java     | 24 ++++++++++---------
 .../python/psi/impl/PyFunctionImpl.java       | 10 ++++----
 .../python/psi/impl/PyNamedParameterImpl.java | 10 ++++----
 4 files changed, 24 insertions(+), 25 deletions(-)

diff --git a/python/psi-api/src/com/jetbrains/python/psi/PyAnnotation.java b/python/psi-api/src/com/jetbrains/python/psi/PyAnnotation.java
index cb37c46cac68..71dc7286d84b 100644
--- a/python/psi-api/src/com/jetbrains/python/psi/PyAnnotation.java
+++ b/python/psi-api/src/com/jetbrains/python/psi/PyAnnotation.java
@@ -22,10 +22,7 @@ import org.jetbrains.annotations.Nullable;
 /**
  * @author yole
  */
-public interface PyAnnotation extends PyElement, StubBasedPsiElement<PyAnnotationStub> {
+public interface PyAnnotation extends PyTypedElement, StubBasedPsiElement<PyAnnotationStub> {
   @Nullable
   PyExpression getValue();
-
-  @Nullable
-  PyClass resolveToClass();
 }
diff --git a/python/src/com/jetbrains/python/psi/impl/PyAnnotationImpl.java b/python/src/com/jetbrains/python/psi/impl/PyAnnotationImpl.java
index 2a95c2bdcbf0..79a226c64eae 100644
--- a/python/src/com/jetbrains/python/psi/impl/PyAnnotationImpl.java
+++ b/python/src/com/jetbrains/python/psi/impl/PyAnnotationImpl.java
@@ -16,14 +16,14 @@
 package com.jetbrains.python.psi.impl;
 
 import com.intellij.lang.ASTNode;
-import com.intellij.psi.PsiElement;
-import com.intellij.psi.PsiPolyVariantReference;
 import com.jetbrains.python.PyElementTypes;
 import com.jetbrains.python.psi.PyAnnotation;
-import com.jetbrains.python.psi.PyClass;
 import com.jetbrains.python.psi.PyExpression;
-import com.jetbrains.python.psi.PyReferenceExpression;
 import com.jetbrains.python.psi.stubs.PyAnnotationStub;
+import com.jetbrains.python.psi.types.PyClassLikeType;
+import com.jetbrains.python.psi.types.PyType;
+import com.jetbrains.python.psi.types.TypeEvalContext;
+import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
 /**
@@ -46,13 +46,15 @@ public class PyAnnotationImpl extends PyBaseElementImpl<PyAnnotationStub> implem
 
   @Nullable
   @Override
-  public PyClass resolveToClass() {
-    PyExpression expr = getValue();
-    if (expr instanceof PyReferenceExpression) {
-      final PsiPolyVariantReference reference = ((PyReferenceExpression)expr).getReference();
-      final PsiElement result = reference.resolve();
-      if (result instanceof PyClass) {
-        return (PyClass) result;
+  public PyType getType(@NotNull TypeEvalContext context, @NotNull TypeEvalContext.Key key) {
+    final PyExpression value = getValue();
+    if (value != null) {
+      final PyType type = context.getType(value);
+      if (type instanceof PyClassLikeType) {
+        final PyClassLikeType classType = (PyClassLikeType)type;
+        if (classType.isDefinition()) {
+          return classType.toInstance();
+        }
       }
     }
     return null;
diff --git a/python/src/com/jetbrains/python/psi/impl/PyFunctionImpl.java b/python/src/com/jetbrains/python/psi/impl/PyFunctionImpl.java
index 246b58d963c3..d7c29fe4a830 100644
--- a/python/src/com/jetbrains/python/psi/impl/PyFunctionImpl.java
+++ b/python/src/com/jetbrains/python/psi/impl/PyFunctionImpl.java
@@ -184,11 +184,11 @@ public class PyFunctionImpl extends PyPresentableElementImpl<PyFunctionStub> imp
       }
     }
     if (context.maySwitchToAST(this) && LanguageLevel.forElement(this).isAtLeast(LanguageLevel.PYTHON30)) {
-      PyAnnotation anno = getAnnotation();
-      if (anno != null) {
-        PyClass pyClass = anno.resolveToClass();
-        if (pyClass != null) {
-          return new PyClassTypeImpl(pyClass, false);
+      final PyAnnotation annotation = getAnnotation();
+      if (annotation != null) {
+        final PyType type = context.getType(annotation);
+        if (type != null) {
+          return type;
         }
       }
     }
diff --git a/python/src/com/jetbrains/python/psi/impl/PyNamedParameterImpl.java b/python/src/com/jetbrains/python/psi/impl/PyNamedParameterImpl.java
index bf792fef0dae..cc9178c07f9e 100644
--- a/python/src/com/jetbrains/python/psi/impl/PyNamedParameterImpl.java
+++ b/python/src/com/jetbrains/python/psi/impl/PyNamedParameterImpl.java
@@ -181,11 +181,11 @@ public class PyNamedParameterImpl extends PyPresentableElementImpl<PyNamedParame
       PyParameterList parameterList = (PyParameterList)parent;
       PyFunction func = parameterList.getContainingFunction();
       if (func != null) {
-        PyAnnotation anno = getAnnotation();
-        if (anno != null) {
-          final PyClass pyClass = anno.resolveToClass();
-          if (pyClass != null) {
-            return new PyClassTypeImpl(pyClass, false);
+        final PyAnnotation annotation = getAnnotation();
+        if (annotation != null) {
+          final PyType type = context.getType(annotation);
+          if (type != null) {
+            return type;
           }
         }
         StructuredDocString docString = func.getStructuredDocString();