support class decorators in Python 2.6 (PY-320)

python/src/com/jetbrains/python/parsing/FunctionParsing.java

@@ -1,28 +1,13 @@
-/*
- *  Copyright 2005 Pythonid Project
- *
- *  Licensed under the Apache License, Version 2.0 (the "License");
- *  you may not use this file except in compliance with the License.
- *  You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- *  Unless required by applicable law or agreed to in writing, software
- *  distributed under the License is distributed on an "AS IS"; BASIS,
- *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- *  See the License for the specific language governing permissions and
- *  limitations under the License.
- */
-
 package com.jetbrains.python.parsing;
 
 import com.intellij.lang.PsiBuilder;
 import com.intellij.openapi.diagnostic.Logger;
 import com.intellij.psi.tree.IElementType;
-import static com.jetbrains.python.PyBundle.message;
 import com.jetbrains.python.PyElementTypes;
 import com.jetbrains.python.PyTokenTypes;
 
+import static com.jetbrains.python.PyBundle.message;
+
 /**
  * @author yole
  */
@@ -53,9 +38,9 @@ public class FunctionParsing extends Parsing {
     getStatementParser().parseSuite(functionMarker, PyElementTypes.FUNCTION_DECLARATION);
   }
 
-  public void parseDecoratedFunctionDeclaration() {
+  public void parseDecoratedDeclaration() {
     assertCurrentToken(PyTokenTypes.AT); // ??? need this?
-    final PsiBuilder.Marker functionMarker = myBuilder.mark();
+    final PsiBuilder.Marker decoratorStartMarker = myBuilder.mark();
     final PsiBuilder.Marker decoListMarker = myBuilder.mark();
     boolean decorated = false;
     while (myBuilder.getTokenType() == PyTokenTypes.AT) {
@@ -75,13 +60,16 @@ public class FunctionParsing extends Parsing {
     if (decorated) decoListMarker.done(PyElementTypes.DECORATOR_LIST);
     //else decoListMarker.rollbackTo(); 
     if (myBuilder.getTokenType() == PyTokenTypes.DEF_KEYWORD) {
-      parseFunctionInnards(functionMarker); // it calls functionMarker.done()
+      parseFunctionInnards(decoratorStartMarker); // it calls decoratorStartMarker.done()
+    }
+    else if (myBuilder.getTokenType() == PyTokenTypes.CLASS_KEYWORD) {
+      getStatementParser().parseClassDeclaration(decoratorStartMarker);
     }
     else {
       myBuilder.error(message("PARSE.expected.@.or.def"));
       PsiBuilder.Marker parameterList = myBuilder.mark(); // To have non-empty parameters list at all the time.
       parameterList.done(PyElementTypes.PARAMETER_LIST);
-      functionMarker.done(PyElementTypes.FUNCTION_DECLARATION);
+      decoratorStartMarker.done(PyElementTypes.FUNCTION_DECLARATION);
     }
   }
 

python/src/com/jetbrains/python/parsing/StatementParsing.java

@@ -73,7 +73,7 @@ public class StatementParsing extends Parsing implements ITokenTypeRemapper {
       return;
     }
     if (firstToken == PyTokenTypes.AT) {
-      getFunctionParser().parseDecoratedFunctionDeclaration();
+      getFunctionParser().parseDecoratedDeclaration();
       return;
     }
     if (firstToken == PyTokenTypes.CLASS_KEYWORD) {
@@ -653,8 +653,12 @@ public class StatementParsing extends Parsing implements ITokenTypeRemapper {
   }
 
   private void parseClassDeclaration() {
-    assertCurrentToken(PyTokenTypes.CLASS_KEYWORD);
     final PsiBuilder.Marker classMarker = myBuilder.mark();
+    parseClassDeclaration(classMarker);
+  }
+
+  public void parseClassDeclaration(PsiBuilder.Marker classMarker) {
+    assertCurrentToken(PyTokenTypes.CLASS_KEYWORD);
     myBuilder.advanceLexer();
     checkMatches(PyTokenTypes.IDENTIFIER, "identifier expected");
     final PsiBuilder.Marker inheritMarker = myBuilder.mark();

python/src/com/jetbrains/python/psi/PyClass.java

@@ -1,19 +1,3 @@
-/*
- *  Copyright 2005 Pythonid Project
- *
- *  Licensed under the Apache License, Version 2.0 (the "License");
- *  you may not use this file except in compliance with the License.
- *  You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- *  Unless required by applicable law or agreed to in writing, software
- *  distributed under the License is distributed on an "AS IS"; BASIS,
- *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- *  See the License for the specific language governing permissions and
- *  limitations under the License.
- */
-
 package com.jetbrains.python.psi;
 
 import com.intellij.psi.PsiElement;
@@ -72,4 +56,7 @@ public interface PyClass extends PsiNamedElement, PyStatement, NameDefiner, PyDo
    * @return True iff this and parent are the same or parent is one of our superclasses.
    */
   boolean isSublclass(PyClass parent);
+
+  @Nullable
+  PyDecoratorList getDecoratorList();
 }

python/src/com/jetbrains/python/psi/PyFileElementType.java

@@ -24,7 +24,7 @@ public class PyFileElementType extends IStubFileElementType {
 
   @Override
   public int getStubVersion() {
-    return 5;
+    return 6;
   }
 
   @Override

python/src/com/jetbrains/python/psi/PyFunction.java

@@ -1,19 +1,3 @@
-/*
- *  Copyright 2005 Pythonid Project
- *
- *  Licensed under the Apache License, Version 2.0 (the "License");
- *  you may not use this file except in compliance with the License.
- *  You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- *  Unless required by applicable law or agreed to in writing, software
- *  distributed under the License is distributed on an "AS IS"; BASIS,
- *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- *  See the License for the specific language governing permissions and
- *  limitations under the License.
- */
-
 package com.jetbrains.python.psi;
 
 import com.intellij.lang.ASTNode;
@@ -25,10 +9,8 @@ import org.jetbrains.annotations.Nullable;
 
 /**
  * Function declaration in source (the <code>def</code> and everything within).
- * User: yole
- * Date: 29.05.2005
- * Time: 23:01:03
- * To change this template use File | Settings | File Templates.
+ *
+ * @author yole
  */
 public interface PyFunction extends PsiNamedElement, PyStatement, NameDefiner, PyDocStringOwner, StubBasedPsiElement<PyFunctionStub> {
   PyFunction[] EMPTY_ARRAY = new PyFunction[0];

python/src/com/jetbrains/python/psi/impl/PyClassImpl.java

@@ -1,19 +1,3 @@
-/*
- *  Copyright 2005 Pythonid Project
- *
- *  Licensed under the Apache License, Version 2.0 (the "License");
- *  you may not use this file except in compliance with the License.
- *  You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- *  Unless required by applicable law or agreed to in writing, software
- *  distributed under the License is distributed on an "AS IS"; BASIS,
- *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- *  See the License for the specific language governing permissions and
- *  limitations under the License.
- */
-
 package com.jetbrains.python.psi.impl;
 
 import com.intellij.lang.ASTNode;
@@ -40,11 +24,7 @@ import javax.swing.*;
 import java.util.*;
 
 /**
- * Created by IntelliJ IDEA.
- * User: yole
- * Date: 03.06.2005
- * Time: 0:27:33
- * To change this template use File | Settings | File Templates.
+ * @author yole
  */
 public class PyClassImpl extends PyPresentableElementImpl<PyClassStub> implements PyClass {
 
@@ -188,6 +168,10 @@ public class PyClassImpl extends PyPresentableElementImpl<PyClassStub> implement
     return false;
   }
 
+  public PyDecoratorList getDecoratorList() {
+    return childToPsi(this, PyElementTypes.DECORATOR_LIST);
+  }
+
   protected List<PyClass> getSuperClassesList() {
     PsiElement[] superClassElements = getSuperClassElements();
     if (superClassElements != null) {

python/testData/psi/ClassDecorators.py

@@ -0,0 +1,4 @@
+@foo
+@bar
+class A:
+  pass

python/testData/psi/ClassDecorators.txt

@@ -0,0 +1,27 @@
+PyFile:ClassDecorators.py
+  PyClass: A
+    PyDecoratorList
+      PyDecorator: @foo
+        PsiElement(Py:AT)('@')
+        PyReferenceExpression: foo
+          PsiElement(Py:IDENTIFIER)('foo')
+        PyArgumentList
+          <empty list>
+      PsiWhiteSpace('\n')
+      PyDecorator: @bar
+        PsiElement(Py:AT)('@')
+        PyReferenceExpression: bar
+          PsiElement(Py:IDENTIFIER)('bar')
+        PyArgumentList
+          <empty list>
+    PsiWhiteSpace('\n')
+    PsiElement(Py:CLASS_KEYWORD)('class')
+    PsiWhiteSpace(' ')
+    PsiElement(Py:IDENTIFIER)('A')
+    PyParenthesizedExpression
+      <empty list>
+    PsiElement(Py:COLON)(':')
+    PsiWhiteSpace('\n  ')
+    PyStatementList
+      PyPassStatement
+        PsiElement(Py:PASS_KEYWORD)('pass')

python/testSrc/com/jetbrains/python/PythonParsingTest.java

@@ -104,6 +104,10 @@ public class PythonParsingTest extends ParsingTestCase {
     doTest(LanguageLevel.PYTHON26);
   }
 
+  public void testClassDecorators() throws Exception {
+    doTest(LanguageLevel.PYTHON26);
+  }
+
   public void doTest() throws Exception {
     doTest(LanguageLevel.PYTHON25);
   }