From ba7d99f47853129150cc1700b8eb8096c3e246aa Mon Sep 17 00:00:00 2001
From: Andrey Vlasovskikh <andrey.vlasovskikh@jetbrains.com>
Date: Tue, 31 Mar 2015 22:25:34 +0300
Subject: [PATCH] Moved annotation-based type inference to PyTypingTypeProvider

We couldn't do check for 'typing.Any' inside PyTypingTypeProvider, so
we had to put some 'typing' knowlegde into PyFunctionImpl and
PyParameterImpl which wasn't elegant. Also we couldn't combine
type information from annotations and default arguments that was
required for Optional[T] for 'x: T = None'. Now we can implement it.
---
 .../jetbrains/python/psi/PyAnnotation.java    |  2 +-
 .../codeInsight/PyTypingTypeProvider.java     | 23 ++++++++++++++++++
 .../python/psi/impl/PyAnnotationImpl.java     | 24 -------------------
 .../python/psi/impl/PyFunctionImpl.java       | 13 ----------
 .../python/psi/impl/PyNamedParameterImpl.java | 11 ---------
 5 files changed, 24 insertions(+), 49 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 804295cb2551..4ce5862310a9 100644
--- a/python/psi-api/src/com/jetbrains/python/psi/PyAnnotation.java
+++ b/python/psi-api/src/com/jetbrains/python/psi/PyAnnotation.java
@@ -22,7 +22,7 @@ import org.jetbrains.annotations.Nullable;
 /**
  * @author yole
  */
-public interface PyAnnotation extends PyTypedElement, StubBasedPsiElement<PyAnnotationStub> {
+public interface PyAnnotation extends PyElement, StubBasedPsiElement<PyAnnotationStub> {
   @Nullable
   PyExpression getValue();
 }
diff --git a/python/src/com/jetbrains/python/codeInsight/PyTypingTypeProvider.java b/python/src/com/jetbrains/python/codeInsight/PyTypingTypeProvider.java
index 8b979de2c850..dc4f88e1ebdd 100644
--- a/python/src/com/jetbrains/python/codeInsight/PyTypingTypeProvider.java
+++ b/python/src/com/jetbrains/python/codeInsight/PyTypingTypeProvider.java
@@ -172,6 +172,10 @@ public class PyTypingTypeProvider extends PyTypeProviderBase {
     if (genericType != null) {
       return genericType;
     }
+    final Ref<PyType> classType = getClassType(expression, context);
+    if (classType != null) {
+      return classType.get();
+    }
     final PyType stringBasedType = getStringBasedType(expression, context);
     if (stringBasedType != null) {
       return stringBasedType;
@@ -179,6 +183,25 @@ public class PyTypingTypeProvider extends PyTypeProviderBase {
     return null;
   }
 
+  @Nullable
+  private static Ref<PyType> getClassType(@NotNull PyExpression expression, @NotNull TypeEvalContext context) {
+    final PyType type = context.getType(expression);
+    if (type != null && isAny(type)) {
+      return Ref.create();
+    }
+    if (type instanceof PyClassLikeType) {
+      final PyClassLikeType classType = (PyClassLikeType)type;
+      if (classType.isDefinition()) {
+        final PyType instanceType = classType.toInstance();
+        return Ref.create(instanceType);
+      }
+    }
+    else if (type instanceof PyNoneType) {
+      return Ref.create(type);
+    }
+    return null;
+  }
+
   @Nullable
   private static Ref<PyType> getOptionalType(@NotNull PyExpression expression, @NotNull TypeEvalContext context) {
     if (expression instanceof PySubscriptionExpression) {
diff --git a/python/src/com/jetbrains/python/psi/impl/PyAnnotationImpl.java b/python/src/com/jetbrains/python/psi/impl/PyAnnotationImpl.java
index 675853c2e7ac..1fcb700def0f 100644
--- a/python/src/com/jetbrains/python/psi/impl/PyAnnotationImpl.java
+++ b/python/src/com/jetbrains/python/psi/impl/PyAnnotationImpl.java
@@ -20,11 +20,6 @@ import com.jetbrains.python.PyElementTypes;
 import com.jetbrains.python.psi.PyAnnotation;
 import com.jetbrains.python.psi.PyExpression;
 import com.jetbrains.python.psi.stubs.PyAnnotationStub;
-import com.jetbrains.python.psi.types.PyClassLikeType;
-import com.jetbrains.python.psi.types.PyNoneType;
-import com.jetbrains.python.psi.types.PyType;
-import com.jetbrains.python.psi.types.TypeEvalContext;
-import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
 /**
@@ -44,23 +39,4 @@ public class PyAnnotationImpl extends PyBaseElementImpl<PyAnnotationStub> implem
   public PyExpression getValue() {
     return findChildByClass(PyExpression.class);
   }
-
-  @Nullable
-  @Override
-  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();
-        }
-      }
-      else if (type instanceof PyNoneType) {
-        return type;
-      }
-    }
-    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 9a2de3224e27..5dd6e678ee64 100644
--- a/python/src/com/jetbrains/python/psi/impl/PyFunctionImpl.java
+++ b/python/src/com/jetbrains/python/psi/impl/PyFunctionImpl.java
@@ -38,7 +38,6 @@ import com.intellij.util.PlatformIcons;
 import com.jetbrains.python.PyElementTypes;
 import com.jetbrains.python.PyNames;
 import com.jetbrains.python.PyTokenTypes;
-import com.jetbrains.python.codeInsight.PyTypingTypeProvider;
 import com.jetbrains.python.codeInsight.controlflow.ControlFlowCache;
 import com.jetbrains.python.codeInsight.controlflow.ScopeOwner;
 import com.jetbrains.python.codeInsight.dataflow.scope.ScopeUtil;
@@ -193,18 +192,6 @@ public class PyFunctionImpl extends PyBaseElementImpl<PyFunctionStub> implements
         return returnType;
       }
     }
-    if (context.maySwitchToAST(this) && LanguageLevel.forElement(this).isAtLeast(LanguageLevel.PYTHON30)) {
-      final PyAnnotation annotation = getAnnotation();
-      if (annotation != null) {
-        final PyType type = context.getType(annotation);
-        if (type != null) {
-          if (PyTypingTypeProvider.isAny(type)) {
-            return null;
-          }
-          return type;
-        }
-      }
-    }
     final PyType docStringType = getReturnTypeFromDocString();
     if (docStringType != null) {
       docStringType.assertValid("from docstring");
diff --git a/python/src/com/jetbrains/python/psi/impl/PyNamedParameterImpl.java b/python/src/com/jetbrains/python/psi/impl/PyNamedParameterImpl.java
index e674a372990a..7c89a1c4770c 100644
--- a/python/src/com/jetbrains/python/psi/impl/PyNamedParameterImpl.java
+++ b/python/src/com/jetbrains/python/psi/impl/PyNamedParameterImpl.java
@@ -33,7 +33,6 @@ import com.jetbrains.python.PyElementTypes;
 import com.jetbrains.python.PyNames;
 import com.jetbrains.python.PyTokenTypes;
 import com.jetbrains.python.PythonDialectsTokenSetProvider;
-import com.jetbrains.python.codeInsight.PyTypingTypeProvider;
 import com.jetbrains.python.codeInsight.controlflow.ScopeOwner;
 import com.jetbrains.python.codeInsight.dataflow.scope.ScopeUtil;
 import com.jetbrains.python.psi.*;
@@ -229,16 +228,6 @@ public class PyNamedParameterImpl extends PyBaseElementImpl<PyNamedParameterStub
             return resultRef.get();
           }
         }
-        final PyAnnotation annotation = getAnnotation();
-        if (annotation != null) {
-          final PyType type = context.getType(annotation);
-          if (type != null) {
-            if (PyTypingTypeProvider.isAny(type)) {
-              return null;
-            }
-            return type;
-          }
-        }
         if (context.maySwitchToAST(this)) {
           final PyExpression defaultValue = getDefaultValue();
           if (defaultValue != null) {