From c69668c37eabe3125bc95e08091f2165d7d91644 Mon Sep 17 00:00:00 2001
From: Semyon Proshev <Semyon.Proshev@jetbrains.com>
Date: Thu, 16 Nov 2017 17:46:40 +0300
Subject: [PATCH] Don't forget about operands union while getting slice type
 (PY-26973)

---
 .../psi/impl/PySliceExpressionImpl.java       | 39 ++++++++++++-------
 .../com/jetbrains/python/PyTypeTest.java      | 11 ++++++
 2 files changed, 35 insertions(+), 15 deletions(-)

diff --git a/python/src/com/jetbrains/python/psi/impl/PySliceExpressionImpl.java b/python/src/com/jetbrains/python/psi/impl/PySliceExpressionImpl.java
index 20a27d6ca6bb..7399449fd4a1 100644
--- a/python/src/com/jetbrains/python/psi/impl/PySliceExpressionImpl.java
+++ b/python/src/com/jetbrains/python/psi/impl/PySliceExpressionImpl.java
@@ -17,6 +17,7 @@ package com.jetbrains.python.psi.impl;
 
 import com.intellij.lang.ASTNode;
 import com.intellij.psi.util.PsiTreeUtil;
+import com.intellij.util.containers.ContainerUtil;
 import com.jetbrains.python.PyNames;
 import com.jetbrains.python.PythonDialectsTokenSetProvider;
 import com.jetbrains.python.psi.PyExpression;
@@ -39,21 +40,7 @@ public class PySliceExpressionImpl extends PyElementImpl implements PySliceExpre
   @Override
   public PyType getType(@NotNull TypeEvalContext context, @NotNull TypeEvalContext.Key key) {
     final PyType type = context.getType(getOperand());
-
-    // TODO: Currently we don't evaluate the static range of the slice, so we have to return a generic tuple type without elements
-    if (type instanceof PyTupleType) {
-      return PyBuiltinCache.getInstance(this).getTupleType();
-    }
-
-    if (type instanceof PyCollectionType) {
-      return type;
-    }
-
-    if (type instanceof PyClassType) {
-      return PyUtil.getReturnTypeOfMember(type, PyNames.GETITEM, null, context);
-    }
-
-    return null;
+    return getSliceType(type, context);
   }
 
   @NotNull
@@ -67,4 +54,26 @@ public class PySliceExpressionImpl extends PyElementImpl implements PySliceExpre
   public PySliceItem getSliceItem() {
     return PsiTreeUtil.getChildOfType(this, PySliceItem.class);
   }
+
+  @Nullable
+  private PyType getSliceType(@Nullable PyType operandType, @NotNull TypeEvalContext context) {
+    // TODO: Currently we don't evaluate the static range of the slice, so we have to return a generic tuple type without elements
+    if (operandType instanceof PyTupleType) {
+      return PyBuiltinCache.getInstance(this).getTupleType();
+    }
+
+    if (operandType instanceof PyCollectionType) {
+      return operandType;
+    }
+
+    if (operandType instanceof PyClassType) {
+      return PyUtil.getReturnTypeOfMember(operandType, PyNames.GETITEM, null, context);
+    }
+
+    if (operandType instanceof PyUnionType) {
+      return PyUnionType.union(ContainerUtil.map(((PyUnionType)operandType).getMembers(), member -> getSliceType(member, context)));
+    }
+
+    return null;
+  }
 }
diff --git a/python/testSrc/com/jetbrains/python/PyTypeTest.java b/python/testSrc/com/jetbrains/python/PyTypeTest.java
index 6f30f4072e27..94f420e5a822 100644
--- a/python/testSrc/com/jetbrains/python/PyTypeTest.java
+++ b/python/testSrc/com/jetbrains/python/PyTypeTest.java
@@ -2576,6 +2576,17 @@ public class PyTypeTest extends PyTestCase {
            "expr = A.B.__new__(A.B)");
   }
 
+  // PY-26973
+  public void testSliceOnUnion() {
+    runWithLanguageLevel(
+      LanguageLevel.PYTHON36,
+      () -> doTest("Union[str, Any]",
+                   "from typing import Union\n" +
+                   "myvar: Union[str, int]\n" +
+                   "expr = myvar[0:3]")
+    );
+  }
+
   private static List<TypeEvalContext> getTypeEvalContexts(@NotNull PyExpression element) {
     return ImmutableList.of(TypeEvalContext.codeAnalysis(element.getProject(), element.getContainingFile()).withTracing(),
                             TypeEvalContext.userInitiated(element.getProject(), element.getContainingFile()).withTracing());