diff --git a/python/src/com/jetbrains/python/psi/impl/PyListCompExpressionImpl.java b/python/src/com/jetbrains/python/psi/impl/PyListCompExpressionImpl.java index f5d7e2e8052c..8e34c714795d 100644 --- a/python/src/com/jetbrains/python/psi/impl/PyListCompExpressionImpl.java +++ b/python/src/com/jetbrains/python/psi/impl/PyListCompExpressionImpl.java @@ -1,11 +1,15 @@ package com.jetbrains.python.psi.impl; import com.intellij.lang.ASTNode; +import com.jetbrains.python.psi.PyClass; import com.jetbrains.python.psi.PyElementVisitor; +import com.jetbrains.python.psi.PyExpression; import com.jetbrains.python.psi.PyListCompExpression; +import com.jetbrains.python.psi.types.PyCollectionTypeImpl; import com.jetbrains.python.psi.types.PyType; import com.jetbrains.python.psi.types.TypeEvalContext; import org.jetbrains.annotations.NotNull; +import org.jetbrains.annotations.Nullable; /** * @author yole @@ -20,7 +24,16 @@ public class PyListCompExpressionImpl extends PyComprehensionElementImpl impleme pyVisitor.visitPyListCompExpression(this); } + @Nullable + @Override public PyType getType(@NotNull TypeEvalContext context) { - return PyBuiltinCache.getInstance(this).getListType(); + final PyExpression resultExpr = getResultExpression(); + final PyBuiltinCache cache = PyBuiltinCache.getInstance(this); + final PyClass list = cache.getClass("list"); + if (resultExpr != null && list != null) { + final PyType elementType = resultExpr.getType(context); + return new PyCollectionTypeImpl(list, false, elementType); + } + return cache.getListType(); } } diff --git a/python/testSrc/com/jetbrains/python/PyTypeTest.java b/python/testSrc/com/jetbrains/python/PyTypeTest.java index 69604e626553..8dca1a6f6738 100644 --- a/python/testSrc/com/jetbrains/python/PyTypeTest.java +++ b/python/testSrc/com/jetbrains/python/PyTypeTest.java @@ -441,6 +441,20 @@ public class PyTypeTest extends PyTestCase { "expr = f()\n"); } + // PY-7020 + public void testListComprehensionType() { + final PyExpression expr = parseExpr("expr = [str(x) for x in range(10)]\n"); + final TypeEvalContext context = TypeEvalContext.slow().withTracing(); + final PyType type = expr.getType(context); + assertNotNull(type); + assertInstanceOf(type, PyCollectionType.class); + assertEquals(type.getName(), "list"); + final PyCollectionType collectionType = (PyCollectionType)type; + final PyType elementType = collectionType.getElementType(context); + assertNotNull(elementType); + assertEquals(elementType.getName(), "str"); + } + private PyExpression parseExpr(String text) { myFixture.configureByText(PythonFileType.INSTANCE, text); return myFixture.findElementByText("expr", PyExpression.class);