Patch #1759: Backport of PEP 3129 class decorators

with some help from Georg

1 files changed, 46 insertions, 24 deletions

diff --git a/Python/ast.c b/Python/ast.c
index aa42fe8..7a1b5bc 100644
--- a/Python/ast.c
+++ b/Python/ast.c
@@ -29,6 +29,7 @@ static asdl_seq *ast_for_suite(struct compiling *, const node *);
 static asdl_seq *ast_for_exprlist(struct compiling *, const node *,
                                   expr_context_ty);
 static expr_ty ast_for_testlist(struct compiling *, const node *);
+static stmt_ty ast_for_classdef(struct compiling *, const node *, asdl_seq *);
 static expr_ty ast_for_testlist_gexp(struct compiling *, const node *);
 
 /* Note different signature for ast_for_call */
@@ -828,27 +829,16 @@ ast_for_decorators(struct compiling *c, const node *n)
 }
 
 static stmt_ty
-ast_for_funcdef(struct compiling *c, const node *n)
+ast_for_funcdef(struct compiling *c, const node *n, asdl_seq *decorator_seq)
 {
-    /* funcdef: 'def' [decorators] NAME parameters ':' suite */
+    /* funcdef: 'def' NAME parameters ':' suite */
     identifier name;
     arguments_ty args;
     asdl_seq *body;
-    asdl_seq *decorator_seq = NULL;
-    int name_i;
+    int name_i = 1;
 
     REQ(n, funcdef);
 
-    if (NCH(n) == 6) { /* decorators are present */
-        decorator_seq = ast_for_decorators(c, CHILD(n, 0));
-        if (!decorator_seq)
-            return NULL;
-        name_i = 2;
-    }
-    else {
-        name_i = 1;
-    }
-
     name = NEW_IDENTIFIER(CHILD(n, name_i));
     if (!name)
         return NULL;
@@ -867,6 +857,36 @@ ast_for_funcdef(struct compiling *c, const node *n)
                        n->n_col_offset, c->c_arena);
 }
 
+static stmt_ty
+ast_for_decorated(struct compiling *c, const node *n)
+{
+    /* decorated: decorators (classdef | funcdef) */
+    stmt_ty thing = NULL;
+    asdl_seq *decorator_seq = NULL;
+
+    REQ(n, decorated);
+
+    decorator_seq = ast_for_decorators(c, CHILD(n, 0));
+    if (!decorator_seq)
+      return NULL;
+
+    assert(TYPE(CHILD(n, 1)) == funcdef ||
+	   TYPE(CHILD(n, 1)) == classdef);
+
+    if (TYPE(CHILD(n, 1)) == funcdef) {
+      thing = ast_for_funcdef(c, CHILD(n, 1), decorator_seq);
+    } else if (TYPE(CHILD(n, 1)) == classdef) {
+      thing = ast_for_classdef(c, CHILD(n, 1), decorator_seq);
+    }
+    /* we count the decorators in when talking about the class' or
+       function's line number */
+    if (thing) {
+        thing->lineno = LINENO(n);
+        thing->col_offset = n->n_col_offset;
+    }
+    return thing;
+}
+
 static expr_ty
 ast_for_lambdef(struct compiling *c, const node *n)
 {
@@ -2968,7 +2988,7 @@ ast_for_with_stmt(struct compiling *c, const node *n)
 }
 
 static stmt_ty
-ast_for_classdef(struct compiling *c, const node *n)
+ast_for_classdef(struct compiling *c, const node *n, asdl_seq *decorator_seq)
 {
     /* classdef: 'class' NAME ['(' testlist ')'] ':' suite */
     asdl_seq *bases, *s;
@@ -2984,16 +3004,16 @@ ast_for_classdef(struct compiling *c, const node *n)
         s = ast_for_suite(c, CHILD(n, 3));
         if (!s)
             return NULL;
-        return ClassDef(NEW_IDENTIFIER(CHILD(n, 1)), NULL, s, LINENO(n),
-                        n->n_col_offset, c->c_arena);
+        return ClassDef(NEW_IDENTIFIER(CHILD(n, 1)), NULL, s, decorator_seq,
+                        LINENO(n), n->n_col_offset, c->c_arena);
     }
     /* check for empty base list */
     if (TYPE(CHILD(n,3)) == RPAR) {
         s = ast_for_suite(c, CHILD(n,5));
         if (!s)
                 return NULL;
-        return ClassDef(NEW_IDENTIFIER(CHILD(n, 1)), NULL, s, LINENO(n),
-                        n->n_col_offset, c->c_arena);
+        return ClassDef(NEW_IDENTIFIER(CHILD(n, 1)), NULL, s, decorator_seq,
+                        LINENO(n), n->n_col_offset, c->c_arena);
     }
 
     /* else handle the base class list */
@@ -3004,8 +3024,8 @@ ast_for_classdef(struct compiling *c, const node *n)
     s = ast_for_suite(c, CHILD(n, 6));
     if (!s)
         return NULL;
-    return ClassDef(NEW_IDENTIFIER(CHILD(n, 1)), bases, s, LINENO(n),
-                    n->n_col_offset, c->c_arena);
+    return ClassDef(NEW_IDENTIFIER(CHILD(n, 1)), bases, s, decorator_seq,
+                    LINENO(n), n->n_col_offset, c->c_arena);
 }
 
 static stmt_ty
@@ -3054,7 +3074,7 @@ ast_for_stmt(struct compiling *c, const node *n)
     }
     else {
         /* compound_stmt: if_stmt | while_stmt | for_stmt | try_stmt
-                        | funcdef | classdef
+                        | funcdef | classdef | decorated
         */
         node *ch = CHILD(n, 0);
         REQ(n, compound_stmt);
@@ -3070,9 +3090,11 @@ ast_for_stmt(struct compiling *c, const node *n)
             case with_stmt:
                 return ast_for_with_stmt(c, ch);
             case funcdef:
-                return ast_for_funcdef(c, ch);
+                return ast_for_funcdef(c, ch, NULL);
             case classdef:
-                return ast_for_classdef(c, ch);
+                return ast_for_classdef(c, ch, NULL);
+	    case decorated:
+	        return ast_for_decorated(c, ch);
             default:
                 PyErr_Format(PyExc_SystemError,
                              "unhandled small_stmt: TYPE=%d NCH=%d\n",