Support decorators on classes as well

src/analysis/type_analysis.cpp

@@ -470,6 +470,14 @@ private:
     }
 
     virtual void visit_classdef(AST_ClassDef* node) {
+        for (auto d : node->decorator_list) {
+            getType(d);
+        }
+
+        for (auto b : node->bases) {
+            getType(b);
+        }
+
         // TODO should we speculate that classdefs will generally return a class?
         // CompilerType* t = typeFromClass(type_cls);
         CompilerType* t = UNKNOWN;
@@ -501,7 +509,17 @@ private:
         }
     }
 
-    virtual void visit_functiondef(AST_FunctionDef* node) { _doSet(node->name, typeFromClass(function_cls)); }
+    virtual void visit_functiondef(AST_FunctionDef* node) {
+        for (auto d : node->decorator_list) {
+            getType(d);
+        }
+
+        for (auto d : node->args->defaults) {
+            getType(d);
+        }
+
+        _doSet(node->name, typeFromClass(function_cls));
+    }
 
     virtual void visit_global(AST_Global* node) {}
 

src/codegen/irgen/irgenerator.cpp

@@ -1290,7 +1290,11 @@ private:
         }
 
         RELEASE_ASSERT(node->bases.size() == 1, "");
-        RELEASE_ASSERT(node->decorator_list.size() == 0, "");
+
+        std::vector<CompilerVariable*> decorators;
+        for (auto d : node->decorator_list) {
+            decorators.push_back(evalExpr(d, exc_info));
+        }
 
         CompilerVariable* base = evalExpr(node->bases[0], exc_info);
         ConcreteCompilerVariable* converted_base = base->makeConverted(emitter, base->getBoxType());
@@ -1324,7 +1328,13 @@ private:
                                   converted_base->getValue(), converted_attr_dict->getValue()).getInstruction();
 
         // Note: createuserClass is free to manufacture non-class objects
-        auto cls = new ConcreteCompilerVariable(UNKNOWN, classobj, true);
+        CompilerVariable* cls = new ConcreteCompilerVariable(UNKNOWN, classobj, true);
+
+        for (int i = decorators.size() - 1; i >= 0; i--) {
+            cls = decorators[i]->call(emitter, getOpInfoForNode(node, exc_info), ArgPassSpec(1), { cls }, NULL);
+            decorators[i]->decvref(emitter);
+        }
+
         _doSet(node->name, cls, exc_info);
         cls->decvref(emitter);
     }
@@ -1503,6 +1513,7 @@ private:
 
         for (int i = decorators.size() - 1; i >= 0; i--) {
             func = decorators[i]->call(emitter, getOpInfoForNode(node, exc_info), ArgPassSpec(1), { func }, NULL);
+            decorators[i]->decvref(emitter);
         }
 
         _doSet(node->name, func, exc_info);

test/tests/classdef_ordering.py

-# expected: fail
-# - decorators
-
-def f(o, msg):
-    print msg
-    return o
-
-@f(lambda c: f(c, "calling decorator"), "evaluating decorator object")
-class C(f(object, "evaluating base")):
-    print "in classdef"

test/tests/decorators.py

@@ -15,6 +15,16 @@ print
 
 
 
+@print_when_eval(lambda f: print_when_eval((f, 2), ("calling outer decorator", f)), "evaluating outer decorator")
+@print_when_eval(lambda f: print_when_eval((f, 1), ("calling inner decorator", f)), "evaluating inner decorator")
+class C(print_when_eval(object, "evaluating base")):
+    print "in classdef"
+    pass
+print
+print C
+print
+
+
 
 
 def print_args(f):