add decorator

Poornima2003 · Feb 25, 2023 · c7bab88 · c7bab88
1 parent 35435ab
commit c7bab88
Show file tree

Hide file tree

Showing 8 changed files with 72 additions and 12 deletions.
diff --git a/src/builtins.h b/src/builtins.h
@@ -394,4 +394,16 @@ def shuffle(L):
 
 def choice(L):
     return L[randint(0, len(L) - 1)]
+)";
+
+const char* kFuncToolsCode = R"(
+def cache(f):
+    def wrapper(*args):
+        if not hasattr(f, 'cache'):
+            f.cache = {}
+        key = args
+        if key not in f.cache:
+            f.cache[key] = f(*args)
+        return f.cache[key]
+    return wrapper
 )";
diff --git a/src/ceval.h b/src/ceval.h
@@ -11,6 +11,7 @@ PyVar VM::run_frame(Frame* frame){
         switch (byte.op)
         {
         case OP_NO_OP: continue;
+        case OP_SETUP_DECORATOR: continue;
         case OP_LOAD_CONST: frame->push(frame->co->consts[byte.arg]); continue;
         case OP_LOAD_FUNCTION: {
             const PyVar obj = frame->co->consts[byte.arg];
@@ -32,13 +33,11 @@ PyVar VM::run_frame(Frame* frame){
             auto& p = frame->co->names[byte.arg];
             NameRef(p).set(this, frame, frame->pop());
         } continue;
-        case OP_BUILD_ATTR: {
-            int name = byte.arg >> 1;
-            bool _rvalue = byte.arg % 2 == 1;
-            auto& attr = frame->co->names[name];
+        case OP_BUILD_ATTR_REF: case OP_BUILD_ATTR: {
+            auto& attr = frame->co->names[byte.arg];
             PyVar obj = frame->pop_value(this);
             AttrRef ref = AttrRef(obj, NameRef(attr));
-            if(_rvalue) frame->push(ref.get(this, frame));
+            if(byte.op == OP_BUILD_ATTR) frame->push(ref.get(this, frame));
             else frame->push(PyRef(ref));
         } continue;
         case OP_BUILD_INDEX: {

diff --git a/src/compiler.h b/src/compiler.h
@@ -206,6 +206,7 @@ class Compiler {
                 case ')': parser->set_next_token(TK(")")); return;
                 case '[': parser->set_next_token(TK("[")); return;
                 case ']': parser->set_next_token(TK("]")); return;
+                case '@': parser->set_next_token(TK("@")); return;
                 case '%': parser->set_next_token_2('=', TK("%"), TK("%=")); return;
                 case '&': parser->set_next_token_2('=', TK("&"), TK("&=")); return;
                 case '|': parser->set_next_token_2('=', TK("|"), TK("|=")); return;
@@ -636,8 +637,7 @@ class Compiler {
         consume(TK("@id"));
         const Str& name = parser->prev.str();
         int index = co()->add_name(name, NAME_ATTR);
-        index = (index<<1) + (int)(co()->_rvalue>0);
-        emit(OP_BUILD_ATTR, index);
+        emit(co()->_rvalue ? OP_BUILD_ATTR : OP_BUILD_ATTR_REF, index);
     }
 
     // [:], [:b]
@@ -746,7 +746,7 @@ class Compiler {
             consume(TK("@id"));
             Token tkname = parser->prev;
             int index = co()->add_name(tkname.str(), NAME_ATTR);
-            emit(OP_BUILD_ATTR, (index<<1)+1);
+            emit(OP_BUILD_ATTR, index);
             if (match(TK("as"))) {
                 consume(TK("@id"));
                 tkname = parser->prev;
@@ -893,6 +893,14 @@ class Compiler {
             compile_from_import();
         } else if (match(TK("def"))){
             compile_function();
+        } else if (match(TK("@"))){
+            EXPR();
+            if(!match_newlines(mode()==REPL_MODE)){
+                SyntaxError("expected a new line after '@'");
+            }
+            emit(OP_SETUP_DECORATOR);
+            consume(TK("def"));
+            compile_function();
         } else if (match(TK("try"))) {
             compile_try_except();
         } else if(match(TK("assert"))) {
@@ -1020,6 +1028,7 @@ class Compiler {
     }
 
     void compile_function(){
+        bool has_decorator = !co()->codes.empty() && co()->codes.back().op == OP_SETUP_DECORATOR;
         if(is_compiling_class){
             if(match(TK("pass"))) return;
             consume(TK("def"));
@@ -1047,14 +1056,19 @@ class Compiler {
         emit(OP_LOAD_FUNCTION, co()->add_const(vm->PyFunction(func)));
         if(name_scope() == NAME_LOCAL) emit(OP_SETUP_CLOSURE);
         if(!is_compiling_class){
-            if(obj_name.empty()) emit(OP_STORE_NAME, co()->add_name(func.name, name_scope()));
-            else {
+            if(obj_name.empty()){
+                if(has_decorator) emit(OP_CALL, 1);
+                emit(OP_STORE_NAME, co()->add_name(func.name, name_scope()));
+            } else {
+                if(has_decorator) SyntaxError("decorator is not supported here");
                 emit(OP_LOAD_NAME, co()->add_name(obj_name, name_scope()));
                 int index = co()->add_name(func.name, NAME_ATTR);
-                emit(OP_BUILD_ATTR, (index<<1)+0);
+                emit(OP_BUILD_ATTR_REF, index);
                 emit(OP_ROT_TWO);
                 emit(OP_STORE_REF);
             }
+        }else{
+            if(has_decorator) SyntaxError("decorator is not supported here");
         }
     }
 

diff --git a/src/obj.h b/src/obj.h
@@ -109,6 +109,8 @@ struct Py_ : PyObject {
             _attr = new pkpy::NameDict(16, kTypeAttrLoadFactor);
         }else if constexpr(std::is_same_v<T, DummyInstance>){
             _attr = new pkpy::NameDict(4, kInstAttrLoadFactor);
+        }else if constexpr(std::is_same_v<T, pkpy::Function> || std::is_same_v<T, pkpy::NativeFunc>){
+            _attr = new pkpy::NameDict(4, kInstAttrLoadFactor);
         }else{
             _attr = nullptr;
         }

diff --git a/src/opcodes.h b/src/opcodes.h
@@ -66,6 +66,7 @@ OPCODE(RE_RAISE)
 
 OPCODE(BUILD_INDEX)
 OPCODE(BUILD_ATTR)
+OPCODE(BUILD_ATTR_REF)
 OPCODE(STORE_NAME)
 OPCODE(STORE_FUNCTION)
 OPCODE(STORE_REF)
@@ -83,5 +84,6 @@ OPCODE(INPLACE_BINARY_OP)
 OPCODE(INPLACE_BITWISE_OP)
 
 OPCODE(SETUP_CLOSURE)
+OPCODE(SETUP_DECORATOR)
 
 #endif
diff --git a/src/parser.h b/src/parser.h
@@ -8,7 +8,7 @@ constexpr const char* kTokens[] = {
     "@error", "@eof", "@eol", "@sof",
     ".", ",", ":", ";", "#", "(", ")", "[", "]", "{", "}", "%", "::",
     "+", "-", "*", "/", "//", "**", "=", ">", "<", "...", "->",
-    "<<", ">>", "&", "|", "^", "?",
+    "<<", ">>", "&", "|", "^", "?", "@",
     "==", "!=", ">=", "<=",
     "+=", "-=", "*=", "/=", "//=", "%=", "&=", "|=", "^=",
     /** KW_BEGIN **/

diff --git a/src/pocketpy.h b/src/pocketpy.h
@@ -70,6 +70,12 @@ void init_builtins(VM* _vm) {
         return vm->new_object(vm->tp_super, *self);
     });
 
+    _vm->bind_builtin_func<1>("id", [](VM* vm, pkpy::Args& args) {
+        const PyVar& obj = args[0];
+        if(obj.is_tagged()) return vm->PyInt((i64)0);
+        return vm->PyInt(obj.bits);
+    });
+
     _vm->bind_builtin_func<1>("eval", [](VM* vm, pkpy::Args& args) {
         CodeObject_ code = vm->compile(vm->PyStr_AS_C(args[0]), "<eval>", EVAL_MODE);
         return vm->_exec(code, vm->top_frame()->_module, vm->top_frame()->_locals);
@@ -776,6 +782,12 @@ void add_module_random(VM* vm){
     vm->_exec(code, mod);
 }
 
+void add_module_functools(VM* vm){
+    PyVar mod = vm->new_module("functools");
+    CodeObject_ code = vm->compile(kFuncToolsCode, "functools.py", EXEC_MODE);
+    vm->_exec(code, mod);
+}
+
 void VM::post_init(){
     init_builtins(this);
     add_module_sys(this);
@@ -787,6 +799,7 @@ void VM::post_init(){
     add_module_random(this);
     add_module_io(this);
     add_module_os(this);
+    add_module_functools(this);
 
     CodeObject_ code = compile(kBuiltinsCode, "<builtins>", EXEC_MODE);
     this->_exec(code, this->builtins);

diff --git a/tests/_decorator.py b/tests/_decorator.py
@@ -0,0 +1,18 @@
+
+def cache(f):
+    def wrapper(*args):
+        if not hasattr(f, 'cache'):
+            f.cache = {}
+        key = args
+        if key not in f.cache:
+            f.cache[key] = f(*args)
+        return f.cache[key]
+    return wrapper
+
+@cache
+def fib(n):
+    if n < 2:
+        return n
+    return fib(n-1) + fib(n-2)
+
+assert fib(32) == 2178309