From a84963397fae514041aa13e2cca969fc60fec12d Mon Sep 17 00:00:00 2001
From: Mikhail Golubev <mikhail.golubev@jetbrains.com>
Date: Fri, 30 Jun 2017 22:26:26 +0300
Subject: [PATCH] PY-18816 Avoid creating code fragment for annotation if AST
 replaced stub already

---
 .../python/codeInsight/typing/PyTypingTypeProvider.java    | 7 +++++--
 python/testSrc/com/jetbrains/python/PyTypingTest.java      | 2 +-
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/python/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java b/python/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java
index 2673c12db103..9bf273caf633 100644
--- a/python/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java
+++ b/python/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java
@@ -22,6 +22,7 @@ import com.google.common.collect.Sets;
 import com.intellij.openapi.util.Ref;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFile;
+import com.intellij.psi.StubBasedPsiElement;
 import com.intellij.psi.impl.source.resolve.FileContextUtil;
 import com.intellij.psi.util.CachedValueProvider;
 import com.intellij.psi.util.CachedValuesManager;
@@ -623,7 +624,8 @@ public class PyTypingTypeProvider extends PyTypeProviderBase {
 
   @Nullable
   private static PyExpression getAnnotationValue(@NotNull PyAnnotationOwner owner, @NotNull TypeEvalContext context) {
-    if (context.maySwitchToAST(owner)) {
+    final boolean elementWasParsed = owner instanceof StubBasedPsiElement && ((StubBasedPsiElement)owner).getStub() == null;
+    if (context.maySwitchToAST(owner) || elementWasParsed) {
       final PyAnnotation annotation = owner.getAnnotation();
       if (annotation != null) {
         return annotation.getValue();
@@ -665,7 +667,8 @@ public class PyTypingTypeProvider extends PyTypeProviderBase {
 
   @Nullable
   private static Ref<PyType> getVariableTypeCommentType(@NotNull String contents, @NotNull PsiElement anchor, @NotNull Context context) {
-    final PyExpression expr = PyUtil.createExpressionFromFragment(contents, anchor);
+    // TODO pass the real anchor as the context element for the fragment to resolve local classes/type aliases
+    final PyExpression expr = PyUtil.createExpressionFromFragment(contents, anchor.getContainingFile());
     if (expr != null) {
       // Such syntax is specific to "# type:" comments, unpacking in type hints is not allowed anywhere else
       if (expr instanceof PyTupleExpression) {
diff --git a/python/testSrc/com/jetbrains/python/PyTypingTest.java b/python/testSrc/com/jetbrains/python/PyTypingTest.java
index 8dfe11baee07..fc168974c6ef 100644
--- a/python/testSrc/com/jetbrains/python/PyTypingTest.java
+++ b/python/testSrc/com/jetbrains/python/PyTypingTest.java
@@ -451,7 +451,7 @@ public class PyTypingTest extends PyTestCase {
 
   // PY-18427
   public void testConditionalType() {
-    doTest("Union[str, int]",
+    doTest("Union[int, str]",
            "if something:\n" +
            "    Type = int\n" +
            "else:\n" +