From e803108eabecfd428cb078b699e2567c6a963c5f Mon Sep 17 00:00:00 2001 From: felfoldy Date: Mon, 6 Jul 2026 19:39:52 +0200 Subject: [PATCH] Add function __annotations__ --- include/pocketpy/objects/codeobject.h | 3 +++ src/compiler/compiler.c | 12 +++++++++-- src/interpreter/vm.c | 1 + src/modules/builtins.c | 12 +++++++++++ src/objects/codeobject.c | 8 +++++++ src/objects/codeobject_ser.c | 31 +++++++++++++++++++++++++-- tests/970_inspect.py | 26 ++++++++++++++++++++++ 7 files changed, 89 insertions(+), 4 deletions(-) diff --git a/include/pocketpy/objects/codeobject.h b/include/pocketpy/objects/codeobject.h index 6691eb68..babc6fe8 100644 --- a/include/pocketpy/objects/codeobject.h +++ b/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); diff --git a/src/compiler/compiler.c b/src/compiler/compiler.c index d7e50cee..b421cf06 100644 --- a/src/compiler/compiler.c +++ b/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)); diff --git a/src/interpreter/vm.c b/src/interpreter/vm.c index bf0636c1..55a4bc67 100644 --- a/src/interpreter/vm.c +++ b/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) { diff --git a/src/modules/builtins.c b/src/modules/builtins.c index af6029c6..99d3980f 100644 --- a/src/modules/builtins.c +++ b/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; } diff --git a/src/objects/codeobject.c b/src/objects/codeobject.c index 3ba251a5..396fd879 100644 --- a/src/objects/codeobject.c +++ b/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); diff --git a/src/objects/codeobject_ser.c b/src/objects/codeobject_ser.c index 236c67ab..e76bc1fb 100644 --- a/src/objects/codeobject_ser.c +++ b/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; } diff --git a/tests/970_inspect.py b/tests/970_inspect.py index aec8ec5c..4137eb70 100644 --- a/tests/970_inspect.py +++ b/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__ == {}