diff --git a/python/src/com/jetbrains/python/parsing/ExpressionParsing.java b/python/src/com/jetbrains/python/parsing/ExpressionParsing.java index 3f5d29ce0336..c9442ece255f 100644 --- a/python/src/com/jetbrains/python/parsing/ExpressionParsing.java +++ b/python/src/com/jetbrains/python/parsing/ExpressionParsing.java @@ -1,5 +1,5 @@ /* - * Copyright 2000-2014 JetBrains s.r.o. + * Copyright 2000-2016 JetBrains s.r.o. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -118,7 +118,7 @@ public class ExpressionParsing extends Parsing { if (!parseSingleExpression(isTargetExpression)) { builder.error(message("PARSE.expected.expression")); } - if (builder.getTokenType() == PyTokenTypes.FOR_KEYWORD) { + if (atForOrAsyncFor()) { parseComprehension(expr, PyTokenTypes.RBRACKET, PyElementTypes.LIST_COMP_EXPRESSION); } else { @@ -153,7 +153,7 @@ public class ExpressionParsing extends Parsing { myBuilder.error(message("PARSE.expected.expression")); } } - if (atToken(PyTokenTypes.FOR_KEYWORD)) { + if (atForOrAsyncFor()) { continue; } if (endToken == null || matchToken(endToken)) { @@ -214,7 +214,7 @@ public class ExpressionParsing extends Parsing { firstExprMarker.drop(); parseSetLiteralTail(expr); } - else if (atToken(PyTokenTypes.FOR_KEYWORD)) { + else if (atForOrAsyncFor()) { firstExprMarker.drop(); parseComprehension(expr, PyTokenTypes.RBRACE, PyElementTypes.SET_COMP_EXPRESSION); } @@ -236,7 +236,7 @@ public class ExpressionParsing extends Parsing { return; } firstKeyValueMarker.done(PyElementTypes.KEY_VALUE_EXPRESSION); - if (myBuilder.getTokenType() == PyTokenTypes.FOR_KEYWORD) { + if (atForOrAsyncFor()) { parseComprehension(startMarker, PyTokenTypes.RBRACE, PyElementTypes.DICT_COMP_EXPRESSION); } else { @@ -299,7 +299,7 @@ public class ExpressionParsing extends Parsing { } else { parseYieldOrTupleExpression(isTargetExpression); - if (myBuilder.getTokenType() == PyTokenTypes.FOR_KEYWORD) { + if (atForOrAsyncFor()) { parseComprehension(expr, PyTokenTypes.RPAR, PyElementTypes.GENERATOR_EXPRESSION); } else { @@ -500,7 +500,7 @@ public class ExpressionParsing extends Parsing { while (myBuilder.getTokenType() != PyTokenTypes.RPAR) { argNumber++; if (argNumber > 1) { - if (argNumber == 2 && atToken(PyTokenTypes.FOR_KEYWORD) && genexpr != null) { + if (argNumber == 2 && atForOrAsyncFor() && genexpr != null) { parseComprehension(genexpr, null, PyElementTypes.GENERATOR_EXPRESSION); genexpr = null; continue; @@ -993,4 +993,18 @@ public class ExpressionParsing extends Parsing { return parseMemberExpression(isTargetExpression); } } + + private boolean atForOrAsyncFor() { + if (atToken(PyTokenTypes.FOR_KEYWORD)) { + return true; + } else if (matchToken(PyTokenTypes.ASYNC_KEYWORD)) { + if (atToken(PyTokenTypes.FOR_KEYWORD)) { + return true; + } else { + myBuilder.error("'for' expected"); + return false; + } + } + return false; + } } diff --git a/python/src/com/jetbrains/python/validation/DumbAwareHighlightingAnnotator.java b/python/src/com/jetbrains/python/validation/DumbAwareHighlightingAnnotator.java index 7ab06c19b322..9fd9a88399a3 100644 --- a/python/src/com/jetbrains/python/validation/DumbAwareHighlightingAnnotator.java +++ b/python/src/com/jetbrains/python/validation/DumbAwareHighlightingAnnotator.java @@ -20,11 +20,13 @@ import com.intellij.lang.ASTNode; import com.intellij.lang.annotation.Annotation; import com.intellij.psi.PsiElement; import com.intellij.psi.PsiFile; +import com.intellij.psi.tree.TokenSet; import com.intellij.psi.util.PsiTreeUtil; import com.jetbrains.python.PyTokenTypes; import com.jetbrains.python.highlighting.PyHighlighter; import com.jetbrains.python.psi.*; import org.jetbrains.annotations.NotNull; +import org.jetbrains.annotations.Nullable; import java.util.Optional; @@ -63,6 +65,8 @@ public class DumbAwareHighlightingAnnotator extends PyAnnotator implements Highl @Override public void visitPyComprehensionElement(PyComprehensionElement node) { + highlightKeywords(node, PyTokenTypes.ASYNC_KEYWORD); + PsiTreeUtil .collectElementsOfType(node.getResultExpression(), PyPrefixExpression.class) .stream() @@ -78,7 +82,16 @@ public class DumbAwareHighlightingAnnotator extends PyAnnotator implements Highl } private void highlightKeyword(@NotNull PsiElement node, @NotNull PyElementType elementType) { - final ASTNode astNode = node.getNode().findChildByType(elementType); + highlightAsKeyword(node.getNode().findChildByType(elementType)); + } + + private void highlightKeywords(@NotNull PsiElement node, @NotNull PyElementType elementType) { + for (ASTNode astNode : node.getNode().getChildren(TokenSet.create(elementType))) { + highlightAsKeyword(astNode); + } + } + + private void highlightAsKeyword(@Nullable ASTNode astNode) { if (astNode != null) { final Annotation annotation = getHolder().createInfoAnnotation(astNode, null); annotation.setTextAttributes(PyHighlighter.PY_KEYWORD); diff --git a/python/testData/psi/AsyncComprehensions.py b/python/testData/psi/AsyncComprehensions.py new file mode 100644 index 000000000000..c318ec551da8 --- /dev/null +++ b/python/testData/psi/AsyncComprehensions.py @@ -0,0 +1,14 @@ +async def async2(): + {i async for i in asyncgen()} + [i async for i in asyncgen()] + {i: i ** 2 async for i in asyncgen()} + (i ** 2 async for i in asyncgen()) + list(i async for i in asyncgen()) + + dataset = {data for line in gen() + async for data in line + if check(data)} + + dataset = {data async for line in asyncgen() + async for data in line + if check(data)} \ No newline at end of file diff --git a/python/testData/psi/AsyncComprehensions.txt b/python/testData/psi/AsyncComprehensions.txt new file mode 100644 index 000000000000..57039334bd92 --- /dev/null +++ b/python/testData/psi/AsyncComprehensions.txt @@ -0,0 +1,248 @@ +PyFile:AsyncComprehensions.py + PyFunction('async2') + PsiElement(Py:ASYNC_KEYWORD)('async') + PsiWhiteSpace(' ') + PsiElement(Py:DEF_KEYWORD)('def') + PsiWhiteSpace(' ') + PsiElement(Py:IDENTIFIER)('async2') + PyParameterList + PsiElement(Py:LPAR)('(') + PsiElement(Py:RPAR)(')') + PsiElement(Py:COLON)(':') + PsiWhiteSpace('\n ') + PyStatementList + PyExpressionStatement + PySetCompExpression + PsiElement(Py:LBRACE)('{') + PyReferenceExpression: i + PsiElement(Py:IDENTIFIER)('i') + PsiWhiteSpace(' ') + PsiElement(Py:ASYNC_KEYWORD)('async') + PsiWhiteSpace(' ') + PsiElement(Py:FOR_KEYWORD)('for') + PsiWhiteSpace(' ') + PyTargetExpression: i + PsiElement(Py:IDENTIFIER)('i') + PsiWhiteSpace(' ') + PsiElement(Py:IN_KEYWORD)('in') + PsiWhiteSpace(' ') + PyCallExpression: asyncgen + PyReferenceExpression: asyncgen + PsiElement(Py:IDENTIFIER)('asyncgen') + PyArgumentList + PsiElement(Py:LPAR)('(') + PsiElement(Py:RPAR)(')') + PsiElement(Py:RBRACE)('}') + PsiWhiteSpace('\n ') + PyExpressionStatement + PyListCompExpression + PsiElement(Py:LBRACKET)('[') + PyReferenceExpression: i + PsiElement(Py:IDENTIFIER)('i') + PsiWhiteSpace(' ') + PsiElement(Py:ASYNC_KEYWORD)('async') + PsiWhiteSpace(' ') + PsiElement(Py:FOR_KEYWORD)('for') + PsiWhiteSpace(' ') + PyTargetExpression: i + PsiElement(Py:IDENTIFIER)('i') + PsiWhiteSpace(' ') + PsiElement(Py:IN_KEYWORD)('in') + PsiWhiteSpace(' ') + PyCallExpression: asyncgen + PyReferenceExpression: asyncgen + PsiElement(Py:IDENTIFIER)('asyncgen') + PyArgumentList + PsiElement(Py:LPAR)('(') + PsiElement(Py:RPAR)(')') + PsiElement(Py:RBRACKET)(']') + PsiWhiteSpace('\n ') + PyExpressionStatement + PyDictCompExpression + PsiElement(Py:LBRACE)('{') + PyKeyValueExpression + PyReferenceExpression: i + PsiElement(Py:IDENTIFIER)('i') + PsiElement(Py:COLON)(':') + PsiWhiteSpace(' ') + PyBinaryExpression + PyReferenceExpression: i + PsiElement(Py:IDENTIFIER)('i') + PsiWhiteSpace(' ') + PsiElement(Py:EXP)('**') + PsiWhiteSpace(' ') + PyNumericLiteralExpression + PsiElement(Py:INTEGER_LITERAL)('2') + PsiWhiteSpace(' ') + PsiElement(Py:ASYNC_KEYWORD)('async') + PsiWhiteSpace(' ') + PsiElement(Py:FOR_KEYWORD)('for') + PsiWhiteSpace(' ') + PyTargetExpression: i + PsiElement(Py:IDENTIFIER)('i') + PsiWhiteSpace(' ') + PsiElement(Py:IN_KEYWORD)('in') + PsiWhiteSpace(' ') + PyCallExpression: asyncgen + PyReferenceExpression: asyncgen + PsiElement(Py:IDENTIFIER)('asyncgen') + PyArgumentList + PsiElement(Py:LPAR)('(') + PsiElement(Py:RPAR)(')') + PsiElement(Py:RBRACE)('}') + PsiWhiteSpace('\n ') + PyExpressionStatement + PyGeneratorExpression + PsiElement(Py:LPAR)('(') + PyBinaryExpression + PyReferenceExpression: i + PsiElement(Py:IDENTIFIER)('i') + PsiWhiteSpace(' ') + PsiElement(Py:EXP)('**') + PsiWhiteSpace(' ') + PyNumericLiteralExpression + PsiElement(Py:INTEGER_LITERAL)('2') + PsiWhiteSpace(' ') + PsiElement(Py:ASYNC_KEYWORD)('async') + PsiWhiteSpace(' ') + PsiElement(Py:FOR_KEYWORD)('for') + PsiWhiteSpace(' ') + PyTargetExpression: i + PsiElement(Py:IDENTIFIER)('i') + PsiWhiteSpace(' ') + PsiElement(Py:IN_KEYWORD)('in') + PsiWhiteSpace(' ') + PyCallExpression: asyncgen + PyReferenceExpression: asyncgen + PsiElement(Py:IDENTIFIER)('asyncgen') + PyArgumentList + PsiElement(Py:LPAR)('(') + PsiElement(Py:RPAR)(')') + PsiElement(Py:RPAR)(')') + PsiWhiteSpace('\n ') + PyExpressionStatement + PyCallExpression: list + PyReferenceExpression: list + PsiElement(Py:IDENTIFIER)('list') + PyArgumentList + PsiElement(Py:LPAR)('(') + PyGeneratorExpression + PyReferenceExpression: i + PsiElement(Py:IDENTIFIER)('i') + PsiWhiteSpace(' ') + PsiElement(Py:ASYNC_KEYWORD)('async') + PsiWhiteSpace(' ') + PsiElement(Py:FOR_KEYWORD)('for') + PsiWhiteSpace(' ') + PyTargetExpression: i + PsiElement(Py:IDENTIFIER)('i') + PsiWhiteSpace(' ') + PsiElement(Py:IN_KEYWORD)('in') + PsiWhiteSpace(' ') + PyCallExpression: asyncgen + PyReferenceExpression: asyncgen + PsiElement(Py:IDENTIFIER)('asyncgen') + PyArgumentList + PsiElement(Py:LPAR)('(') + PsiElement(Py:RPAR)(')') + PsiElement(Py:RPAR)(')') + PsiWhiteSpace('\n\n ') + PyAssignmentStatement + PyTargetExpression: dataset + PsiElement(Py:IDENTIFIER)('dataset') + PsiWhiteSpace(' ') + PsiElement(Py:EQ)('=') + PsiWhiteSpace(' ') + PySetCompExpression + PsiElement(Py:LBRACE)('{') + PyReferenceExpression: data + PsiElement(Py:IDENTIFIER)('data') + PsiWhiteSpace(' ') + PsiElement(Py:FOR_KEYWORD)('for') + PsiWhiteSpace(' ') + PyTargetExpression: line + PsiElement(Py:IDENTIFIER)('line') + PsiWhiteSpace(' ') + PsiElement(Py:IN_KEYWORD)('in') + PsiWhiteSpace(' ') + PyCallExpression: gen + PyReferenceExpression: gen + PsiElement(Py:IDENTIFIER)('gen') + PyArgumentList + PsiElement(Py:LPAR)('(') + PsiElement(Py:RPAR)(')') + PsiWhiteSpace('\n ') + PsiElement(Py:ASYNC_KEYWORD)('async') + PsiWhiteSpace(' ') + PsiElement(Py:FOR_KEYWORD)('for') + PsiWhiteSpace(' ') + PyTargetExpression: data + PsiElement(Py:IDENTIFIER)('data') + PsiWhiteSpace(' ') + PsiElement(Py:IN_KEYWORD)('in') + PsiWhiteSpace(' ') + PyReferenceExpression: line + PsiElement(Py:IDENTIFIER)('line') + PsiWhiteSpace('\n ') + PsiElement(Py:IF_KEYWORD)('if') + PsiWhiteSpace(' ') + PyCallExpression: check + PyReferenceExpression: check + PsiElement(Py:IDENTIFIER)('check') + PyArgumentList + PsiElement(Py:LPAR)('(') + PyReferenceExpression: data + PsiElement(Py:IDENTIFIER)('data') + PsiElement(Py:RPAR)(')') + PsiElement(Py:RBRACE)('}') + PsiWhiteSpace('\n\n ') + PyAssignmentStatement + PyTargetExpression: dataset + PsiElement(Py:IDENTIFIER)('dataset') + PsiWhiteSpace(' ') + PsiElement(Py:EQ)('=') + PsiWhiteSpace(' ') + PySetCompExpression + PsiElement(Py:LBRACE)('{') + PyReferenceExpression: data + PsiElement(Py:IDENTIFIER)('data') + PsiWhiteSpace(' ') + PsiElement(Py:ASYNC_KEYWORD)('async') + PsiWhiteSpace(' ') + PsiElement(Py:FOR_KEYWORD)('for') + PsiWhiteSpace(' ') + PyTargetExpression: line + PsiElement(Py:IDENTIFIER)('line') + PsiWhiteSpace(' ') + PsiElement(Py:IN_KEYWORD)('in') + PsiWhiteSpace(' ') + PyCallExpression: asyncgen + PyReferenceExpression: asyncgen + PsiElement(Py:IDENTIFIER)('asyncgen') + PyArgumentList + PsiElement(Py:LPAR)('(') + PsiElement(Py:RPAR)(')') + PsiWhiteSpace('\n ') + PsiElement(Py:ASYNC_KEYWORD)('async') + PsiWhiteSpace(' ') + PsiElement(Py:FOR_KEYWORD)('for') + PsiWhiteSpace(' ') + PyTargetExpression: data + PsiElement(Py:IDENTIFIER)('data') + PsiWhiteSpace(' ') + PsiElement(Py:IN_KEYWORD)('in') + PsiWhiteSpace(' ') + PyReferenceExpression: line + PsiElement(Py:IDENTIFIER)('line') + PsiWhiteSpace('\n ') + PsiElement(Py:IF_KEYWORD)('if') + PsiWhiteSpace(' ') + PyCallExpression: check + PyReferenceExpression: check + PsiElement(Py:IDENTIFIER)('check') + PyArgumentList + PsiElement(Py:LPAR)('(') + PyReferenceExpression: data + PsiElement(Py:IDENTIFIER)('data') + PsiElement(Py:RPAR)(')') + PsiElement(Py:RBRACE)('}') \ No newline at end of file diff --git a/python/testSrc/com/jetbrains/python/PythonParsingTest.java b/python/testSrc/com/jetbrains/python/PythonParsingTest.java index 605b56ac0d78..92d549edc2d5 100644 --- a/python/testSrc/com/jetbrains/python/PythonParsingTest.java +++ b/python/testSrc/com/jetbrains/python/PythonParsingTest.java @@ -524,6 +524,11 @@ public class PythonParsingTest extends ParsingTestCase { doTest(LanguageLevel.PYTHON36); } + // PY-20770 + public void testAsyncComprehensions() { + doTest(LanguageLevel.PYTHON36); + } + public void doTest(LanguageLevel languageLevel) { LanguageLevel prev = myLanguageLevel; myLanguageLevel = languageLevel;