PY-20744 Parse PEP-526 variable annotations
[idea/community.git] / python / src / com / jetbrains / python / parsing / FunctionParsing.java
1 /*
2  * Copyright 2000-2016 JetBrains s.r.o.
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  * http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 package com.jetbrains.python.parsing;
17
18 import com.intellij.lang.PsiBuilder;
19 import com.intellij.lang.WhitespacesBinders;
20 import com.intellij.psi.tree.IElementType;
21 import com.jetbrains.python.PyElementTypes;
22 import com.jetbrains.python.PyTokenTypes;
23 import org.jetbrains.annotations.NotNull;
24
25 import static com.jetbrains.python.PyBundle.message;
26
27 /**
28  * @author yole
29  */
30 public class FunctionParsing extends Parsing {
31   private static final IElementType FUNCTION_TYPE = PyElementTypes.FUNCTION_DECLARATION;
32
33   public FunctionParsing(ParsingContext context) {
34     super(context);
35   }
36
37   public void parseFunctionDeclaration(@NotNull PsiBuilder.Marker endMarker, boolean async) {
38     assertCurrentToken(PyTokenTypes.DEF_KEYWORD);
39     parseFunctionInnards(endMarker, async);
40   }
41
42   protected IElementType getFunctionType() {
43     return FUNCTION_TYPE;
44   }
45
46   protected void parseFunctionInnards(@NotNull PsiBuilder.Marker functionMarker, boolean async) {
47     myBuilder.advanceLexer();
48     parseIdentifierOrSkip(PyTokenTypes.LPAR);
49     parseParameterList();
50     parseReturnTypeAnnotation();
51     checkMatches(PyTokenTypes.COLON, message("PARSE.expected.colon"));
52     final ParsingContext context = getParsingContext();
53     context.pushScope(context.getScope().withFunction(async));
54     getStatementParser().parseSuite(functionMarker, getFunctionType());
55     context.popScope();
56   }
57
58   public void parseReturnTypeAnnotation() {
59     if (myContext.getLanguageLevel().isPy3K() && myBuilder.getTokenType() == PyTokenTypes.RARROW) {
60       PsiBuilder.Marker maybeReturnAnnotation = myBuilder.mark();
61       nextToken();
62       if (!myContext.getExpressionParser().parseSingleExpression(false)) {
63         myBuilder.error(message("PARSE.expected.expression"));
64       }
65       maybeReturnAnnotation.done(PyElementTypes.ANNOTATION);
66     }
67   }
68
69   public void parseDecoratedDeclaration() {
70     assertCurrentToken(PyTokenTypes.AT); // ??? need this?
71     final PsiBuilder.Marker decoratorStartMarker = myBuilder.mark();
72     final PsiBuilder.Marker decoListMarker = myBuilder.mark();
73     boolean decorated = false;
74     while (myBuilder.getTokenType() == PyTokenTypes.AT) {
75       PsiBuilder.Marker decoratorMarker = myBuilder.mark();
76       myBuilder.advanceLexer();
77       getStatementParser().parseDottedName();
78       if (myBuilder.getTokenType() == PyTokenTypes.LPAR) {
79         getExpressionParser().parseArgumentList();
80       }
81       else { // empty arglist node, so we always have it
82         PsiBuilder.Marker argListMarker = myBuilder.mark();
83         argListMarker.setCustomEdgeTokenBinders(WhitespacesBinders.GREEDY_LEFT_BINDER, null);
84         argListMarker.done(PyElementTypes.ARGUMENT_LIST);
85       }
86       if (atToken(PyTokenTypes.STATEMENT_BREAK)) {
87         decoratorMarker.done(PyElementTypes.DECORATOR_CALL);
88         nextToken();
89       }
90       else {
91         myBuilder.error(message("PARSE.expected.statement.break"));
92         decoratorMarker.done(PyElementTypes.DECORATOR_CALL);
93       }
94       decorated = true;
95     }
96     if (decorated) decoListMarker.done(PyElementTypes.DECORATOR_LIST);
97     //else decoListMarker.rollbackTo();
98     parseDeclarationAfterDecorator(decoratorStartMarker);
99   }
100
101   private void parseDeclarationAfterDecorator(PsiBuilder.Marker endMarker) {
102     if (myBuilder.getTokenType() == PyTokenTypes.ASYNC_KEYWORD) {
103       myBuilder.advanceLexer();
104       parseDeclarationAfterDecorator(endMarker, true);
105     }
106     else {
107       parseDeclarationAfterDecorator(endMarker, false);
108     }
109   }
110
111   protected void parseDeclarationAfterDecorator(PsiBuilder.Marker endMarker, boolean async) {
112     if (myBuilder.getTokenType() == PyTokenTypes.DEF_KEYWORD) {
113       parseFunctionInnards(endMarker, async);
114     }
115     else if (myBuilder.getTokenType() == PyTokenTypes.CLASS_KEYWORD) {
116       getStatementParser().parseClassDeclaration(endMarker);
117     }
118     else {
119       myBuilder.error(message("PARSE.expected.@.or.def"));
120       PsiBuilder.Marker parameterList = myBuilder.mark(); // To have non-empty parameters list at all the time.
121       parameterList.done(PyElementTypes.PARAMETER_LIST);
122       myBuilder.mark().done(PyElementTypes.STATEMENT_LIST); // To have non-empty empty statement list
123       endMarker.done(getFunctionType());
124     }
125   }
126
127   public void parseParameterList() {
128     final PsiBuilder.Marker parameterList;
129     if (myBuilder.getTokenType() != PyTokenTypes.LPAR) {
130       myBuilder.error(message("PARSE.expected.lpar"));
131       parameterList = myBuilder.mark(); // To have non-empty parameters list at all the time.
132       parameterList.done(PyElementTypes.PARAMETER_LIST);
133       return;
134     }
135     parseParameterListContents(PyTokenTypes.RPAR, true, false);
136   }
137
138   public void parseParameterListContents(IElementType endToken, boolean advanceLexer, boolean isLambda) {
139     PsiBuilder.Marker parameterList;
140     parameterList = myBuilder.mark();
141     if (advanceLexer) {
142       myBuilder.advanceLexer();
143     }
144
145     boolean first = true;
146     boolean afterStarParameter = false;
147     while (myBuilder.getTokenType() != endToken) {
148       if (first) {
149         first = false;
150       }
151       else {
152         if (myBuilder.getTokenType() == PyTokenTypes.COMMA) {
153           myBuilder.advanceLexer();
154         }
155         else {
156           myBuilder.error(message("PARSE.expected.comma.lpar.rpar"));
157           break;
158         }
159       }
160       if (myBuilder.getTokenType() == PyTokenTypes.LPAR) {
161         parseParameterSubList();
162         continue;
163       }
164       boolean isStarParameter = atAnyOfTokens(PyTokenTypes.MULT, PyTokenTypes.EXP);
165       if (!parseParameter(endToken, isLambda)) {
166         if (afterStarParameter) {
167           myBuilder.error("expression expected");
168         }
169         break;
170       }
171       if (isStarParameter) {
172         afterStarParameter = true;
173       }
174     }
175
176     if (myBuilder.getTokenType() == endToken && endToken == PyTokenTypes.RPAR) {
177       myBuilder.advanceLexer();
178     }
179
180     parameterList.done(PyElementTypes.PARAMETER_LIST);
181
182     if (myBuilder.getTokenType() == endToken && endToken == PyTokenTypes.COLON) {
183       myBuilder.advanceLexer();
184     }
185   }
186
187   protected boolean parseParameter(IElementType endToken, boolean isLambda) {
188     final PsiBuilder.Marker parameter = myBuilder.mark();
189     boolean isStarParameter = false;
190     if (myBuilder.getTokenType() == PyTokenTypes.MULT) {
191       myBuilder.advanceLexer();
192       if (myContext.getLanguageLevel().isPy3K() &&
193           (myBuilder.getTokenType() == PyTokenTypes.COMMA) || myBuilder.getTokenType() == endToken) {
194         parameter.done(PyElementTypes.SINGLE_STAR_PARAMETER);
195         return true;
196       }
197       isStarParameter = true;
198     }
199     else if (myBuilder.getTokenType() == PyTokenTypes.EXP) {
200       myBuilder.advanceLexer();
201       isStarParameter = true;
202     }
203     if (matchToken(PyTokenTypes.IDENTIFIER)) {
204       if (!isLambda) {
205         parseParameterAnnotation();
206       }
207       if (!isStarParameter && matchToken(PyTokenTypes.EQ)) {
208         if (!getExpressionParser().parseSingleExpression(false)) {
209           PsiBuilder.Marker invalidElements = myBuilder.mark();
210           while(!atAnyOfTokens(endToken, PyTokenTypes.LINE_BREAK, PyTokenTypes.COMMA, null)) {
211             nextToken();
212           }
213           invalidElements.error(message("PARSE.expected.expression"));
214         }
215       }
216       parameter.done(PyElementTypes.NAMED_PARAMETER);
217     }
218     else {
219       parameter.rollbackTo();
220       if (atToken(endToken)) {
221         return false;
222       }
223       PsiBuilder.Marker invalidElements = myBuilder.mark();
224       while (!atToken(endToken) && !atAnyOfTokens(PyTokenTypes.LINE_BREAK, PyTokenTypes.COMMA, null)) {
225         nextToken();
226       }
227       invalidElements.error(message("PARSE.expected.formal.param.name"));
228       return atToken(endToken) || atToken(PyTokenTypes.COMMA);
229     }
230     return true;
231   }
232
233   public void parseParameterAnnotation() {
234     if (myContext.getLanguageLevel().isPy3K() && atToken(PyTokenTypes.COLON)) {
235       PsiBuilder.Marker annotationMarker = myBuilder.mark();
236       nextToken();
237       if (!getExpressionParser().parseSingleExpression(false)) {
238         myBuilder.error(message("PARSE.expected.expression"));
239       }
240       annotationMarker.done(PyElementTypes.ANNOTATION);
241     }
242   }
243
244   private void parseParameterSubList() {
245     assertCurrentToken(PyTokenTypes.LPAR);
246     final PsiBuilder.Marker tuple = myBuilder.mark();
247     myBuilder.advanceLexer();
248     while (true) {
249       if (myBuilder.getTokenType() == PyTokenTypes.IDENTIFIER) {
250         final PsiBuilder.Marker parameter = myBuilder.mark();
251         myBuilder.advanceLexer();
252         parameter.done(PyElementTypes.NAMED_PARAMETER);
253       }
254       else if (myBuilder.getTokenType() == PyTokenTypes.LPAR) {
255         parseParameterSubList();
256       }
257       if (myBuilder.getTokenType() == PyTokenTypes.RPAR) {
258         myBuilder.advanceLexer();
259         break;
260       }
261       if (myBuilder.getTokenType() != PyTokenTypes.COMMA) {
262         myBuilder.error(message("PARSE.expected.comma.lpar.rpar"));
263         break;
264       }
265       myBuilder.advanceLexer();
266     }
267     if (myBuilder.getTokenType() == PyTokenTypes.EQ) {
268       myBuilder.advanceLexer();
269       getExpressionParser().parseSingleExpression(false);
270     }
271     tuple.done(PyElementTypes.TUPLE_PARAMETER);
272   }
273 }