diff --git a/src/compiler.c b/src/compiler.c index 8c264c9..f0601d2 100644 --- a/src/compiler.c +++ b/src/compiler.c @@ -188,7 +188,7 @@ static void declaration(); static void or_(int canAssign); static void ternary(int canAssign); static void and_(int canAssign); -static void classDeclaration(); +static KrkToken classDeclaration(); static void declareVariable(); static void namedVariable(KrkToken name, int canAssign); static void addLocal(KrkToken name); @@ -761,7 +761,11 @@ static void declaration() { } else if (match(TOKEN_LET)) { letDeclaration(); } else if (check(TOKEN_CLASS)) { - classDeclaration(); + KrkToken className = classDeclaration(); + size_t classConst = identifierConstant(&className); + parser.previous = className; + declareVariable(); + defineVariable(classConst); } else if (check(TOKEN_AT)) { decorator(0, TYPE_FUNCTION); } else if (match(TOKEN_EOL) || match(TOKEN_EOF)) { @@ -1018,11 +1022,17 @@ static void method(size_t blockWidth) { } } -static void classDeclaration() { +static KrkToken classDeclaration() { size_t blockWidth = (parser.previous.type == TOKEN_INDENTATION) ? parser.previous.length : 0; advance(); /* Collect the `class` */ consume(TOKEN_IDENTIFIER, "Expected class name."); + Compiler subcompiler; + initCompiler(&subcompiler, TYPE_LAMBDA); + subcompiler.function->chunk.filename = subcompiler.enclosing->function->chunk.filename; + + beginScope(); + KrkToken className = parser.previous; size_t constInd = identifierConstant(&parser.previous); declareVariable(); @@ -1094,8 +1104,15 @@ static void classDeclaration() { } /* else empty class (and at end of file?) we'll allow it for now... */ _pop_class: emitByte(OP_FINALIZE); - endScope(); currentClass = currentClass->enclosing; + KrkFunction * makeclass = endCompiler(); + size_t indFunc = krk_addConstant(currentChunk(), OBJECT_VAL(makeclass)); + EMIT_CONSTANT_OP(OP_CLOSURE, indFunc); + doUpvalues(&subcompiler, makeclass); + freeCompiler(&subcompiler); + emitBytes(OP_CALL, 0); + + return className; } static void markInitialized() { @@ -1187,6 +1204,12 @@ static KrkToken decorator(size_t level, FunctionType type) { function(type, blockWidth); } else if (check(TOKEN_AT)) { funcName = decorator(level+1, type); + } else if (check(TOKEN_CLASS)) { + if (type != TYPE_FUNCTION) { + error("Invalid decorator applied to class"); + return funcName; + } + funcName = classDeclaration(); } else { error("Expected a function declaration or another decorator."); return funcName; diff --git a/src/vm.c b/src/vm.c index 7996477..cc6662d 100644 --- a/src/vm.c +++ b/src/vm.c @@ -2066,7 +2066,6 @@ static KrkValue run() { KrkClass * _class = AS_CLASS(krk_peek(0)); /* Store special methods for quick access */ krk_finalizeClass(_class); - krk_pop(); /* Pop the class as we're done attaching methods */ break; } case OP_INHERIT: { diff --git a/test/testClassDecorator.krk b/test/testClassDecorator.krk new file mode 100644 index 0000000..5bf1f83 --- /dev/null +++ b/test/testClassDecorator.krk @@ -0,0 +1,18 @@ +def classDecorator(cls): + class NewClass(cls): + def extraMethod(self): + print("decorator-added method",self.__class__.__name__) + return NewClass + + +@classDecorator +class FooBar(object): + def __init__(self): + self.bar = "baz" + def myMethod(self): + print(self.bar) + + +let l = FooBar() +l.extraMethod() +l.myMethod() diff --git a/test/testClassDecorator.krk.expect b/test/testClassDecorator.krk.expect new file mode 100644 index 0000000..90b72a1 --- /dev/null +++ b/test/testClassDecorator.krk.expect @@ -0,0 +1,2 @@ +decorator-added method NewClass +baz