Add function __annotations__

include/pocketpy/objects/codeobject.h

@@ -114,6 +114,8 @@ typedef struct FuncDecl {
     int starred_kwarg;  // index in co->varnames, -1 if no **kwarg
     bool nested;        // whether this function is nested
 
+    py_TValue annotations;  // dict[str, str], nil if empty
+
     char* docstring;
 
     FuncType type;
@@ -128,6 +130,7 @@ void FuncDecl__add_arg(FuncDecl* self, py_Name name);
 void FuncDecl__add_kwarg(FuncDecl* self, py_Name name, const py_TValue* value);
 void FuncDecl__add_starred_arg(FuncDecl* self, py_Name name);
 void FuncDecl__add_starred_kwarg(FuncDecl* self, py_Name name);
+void FuncDecl__add_annotation(FuncDecl* self, py_Name name, c11_sv hint);
 void FuncDecl__gc_mark(const FuncDecl* self, c11_vector* p_stack);
 void FuncDecl__dtor(FuncDecl* self);
 

src/compiler/compiler.c

@@ -2368,7 +2368,11 @@ static Error* _compile_f_args(Compiler* self, FuncDecl* decl, bool is_lambda) {
         }
 
         // eat type hints
-        if(!is_lambda && match(TK_COLON)) check(consume_type_hints(self));
+        if(!is_lambda && match(TK_COLON)) {
+            c11_sv hint;
+            check(consume_type_hints_sv(self, &hint));
+            FuncDecl__add_annotation(decl, name, hint);
+        }
         if(state == 0 && curr()->type == TK_ASSIGN) state = 2;
         switch(state) {
             case 0: FuncDecl__add_arg(decl, name); break;
@@ -2417,7 +2421,11 @@ static Error* compile_function(Compiler* self, int decorators) {
         check(_compile_f_args(self, decl, false));
         consume(TK_RPAREN);
     }
-    if(match(TK_ARROW)) check(consume_type_hints(self));
+    if(match(TK_ARROW)) {
+        c11_sv hint;
+        check(consume_type_hints_sv(self, &hint));
+        FuncDecl__add_annotation(decl, py_name("return"), hint);
+    }
     check(compile_block_body(self));
     check(pop_context(self));
 

src/interpreter/vm.c

@@ -635,6 +635,7 @@ void FuncDecl__gc_mark(const FuncDecl* self, c11_vector* p_stack) {
         FuncDeclKwArg* kw = c11__at(FuncDeclKwArg, &self->kwargs, j);
         pk__mark_value(&kw->value);
     }
+    pk__mark_value(&self->annotations);
 }
 
 void CodeObject__gc_mark(const CodeObject* self, c11_vector* p_stack) {

src/modules/builtins.c

@@ -562,6 +562,17 @@ static bool function__doc__(int argc, py_Ref argv) {
     return true;
 }
 
+static bool function__annotations__(int argc, py_Ref argv) {
+    PY_CHECK_ARGC(1);
+    Function* func = py_touserdata(py_arg(0));
+    if(py_isnil(&func->decl->annotations)) {
+        py_newdict(py_retval());
+    } else {
+        py_assign(py_retval(), &func->decl->annotations);
+    }
+    return true;
+}
+
 static bool function__name__(int argc, py_Ref argv) {
     PY_CHECK_ARGC(1);
     Function* func = py_touserdata(py_arg(0));
@@ -589,6 +600,7 @@ py_Type pk_function__register() {
         pk_newtype("function", tp_object, NULL, (void (*)(void*))Function__dtor, false, true);
     py_bindproperty(type, "__doc__", function__doc__, NULL);
     py_bindproperty(type, "__name__", function__name__, NULL);
+    py_bindproperty(type, "__annotations__", function__annotations__, NULL);
     py_bindmagic(type, __repr__, function__repr__);
     return type;
 }

src/objects/codeobject.c

@@ -37,6 +37,7 @@ FuncDecl_ FuncDecl__rcnew(SourceData_ src, c11_sv name) {
     self->starred_arg = -1;
     self->starred_kwarg = -1;
     self->nested = false;
+    self->annotations = *py_NIL();
 
     self->docstring = NULL;
     self->type = FuncType_UNSET;
@@ -90,6 +91,13 @@ void FuncDecl__add_starred_kwarg(FuncDecl* self, py_Name name) {
     self->starred_kwarg = index;
 }
 
+void FuncDecl__add_annotation(FuncDecl* self, py_Name name, c11_sv hint) {
+    if(py_isnil(&self->annotations)) py_newdict(&self->annotations);
+    py_TValue value;
+    py_newstrv(&value, hint);
+    py_dict_setitem_by_str(&self->annotations, py_name2str(name), &value);
+}
+
 void CodeObject__ctor(CodeObject* self, SourceData_ src, c11_sv name) {
     self->src = src;
     PK_INCREF(src);

src/objects/codeobject_ser.c

@@ -5,8 +5,8 @@
 // Magic number for CodeObject serialization: "CO" = 0x434F
 #define CODEOBJECT_MAGIC 0x434F
 #define CODEOBJECT_VER_MAJOR 1
-#define CODEOBJECT_VER_MINOR 0
-#define CODEOBJECT_VER_MINOR_MIN 0
+#define CODEOBJECT_VER_MINOR 1
+#define CODEOBJECT_VER_MINOR_MIN 1
 
 // Forward declarations
 static void FuncDecl__serialize(c11_serializer* s,
@@ -277,6 +277,13 @@ static CodeObject CodeObject__deserialize(c11_deserializer* d, const char* filen
     return co;
 }
 
+static bool annotation__serialize(py_Ref key, py_Ref val, void* ctx) {
+    c11_serializer* s = ctx;
+    c11_serializer__write_cstr(s, py_tostr(key));
+    c11_serializer__write_cstr(s, py_tostr(val));
+    return true;
+}
+
 // Serialize FuncDecl
 static void FuncDecl__serialize(c11_serializer* s,
                                 const FuncDecl* decl,
@@ -317,6 +324,14 @@ static void FuncDecl__serialize(c11_serializer* s,
 
     // type
     c11_serializer__write_i8(s, (int8_t)decl->type);
+
+    // annotations
+    py_Ref annotations = (py_Ref)&decl->annotations;
+    bool empty = py_isnil(annotations);
+    c11_serializer__write_i32(s, empty ? 0 : py_dict_len(annotations));
+    c11_serializer__write_mark(s, '[');
+    if(!empty) py_dict_apply(annotations, annotation__serialize, s);
+    c11_serializer__write_mark(s, ']');
 }
 
 // Deserialize FuncDecl
@@ -374,6 +389,18 @@ static FuncDecl_ FuncDecl__deserialize(c11_deserializer* d, SourceData_ embedded
 
     // type
     self->type = (FuncType)c11_deserializer__read_i8(d);
+
+    // annotations
+    self->annotations = *py_NIL();
+    int annotations_len = c11_deserializer__read_i32(d);
+    c11_deserializer__consume_mark(d, '[');
+    for(int i = 0; i < annotations_len; i++) {
+        const char* name_str = c11_deserializer__read_cstr(d);
+        const char* hint_str = c11_deserializer__read_cstr(d);
+        c11_sv hint = {hint_str, (int)strlen(hint_str)};
+        FuncDecl__add_annotation(self, py_name(name_str), hint);
+    }
+    c11_deserializer__consume_mark(d, ']');
     return self;
 }
 

tests/970_inspect.py

@@ -180,3 +180,29 @@ try:
     exit(1)
 except TypeError:
     pass
+
+# ---------------- __annotations__ ----------------
+def h1(a: int, *args: int, b: str = 'x', c: float = 1.5, **kwargs: str) -> bool:
+    pass
+
+assert h1.__annotations__ == {
+    'a': 'int',
+    'args': 'int',
+    'b': 'str',
+    'c': 'float',
+    'kwargs': 'str',
+    'return': 'bool',
+}
+
+def h2(a, b: int, c=1):
+    pass
+
+assert h2.__annotations__ == {'b': 'int'}
+
+# complex annotation expressions are preserved as written
+def h3(p: list[int], q: dict[str, int]) -> 'A | None':
+    pass
+
+assert h3.__annotations__ == {'p': 'list[int]', 'q': 'dict[str, int]', 'return': "'A | None'"}
+
+assert (lambda x: x).__annotations__ == {}