Merge 56ba5f7f832b929bf98f5ef2f53c106501418dab into acc4f4104bb9e6a5ad5f34a4f9c8e2f354cec20c

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);
 

python/inspect.py

@@ -11,15 +11,20 @@ class Parameter:
 
     empty = _empty
 
-    def __init__(self, name, kind, *default):
+    def __init__(self, name, kind, *default, annotation=None):
         self.name = name
         self.kind = kind
         # pocketpy only allows literal defaults, so use *default as sentinel
         self.default = default[0] if default else _empty
+        self.annotation = _empty if annotation is None else annotation
 
     def __str__(self):
         res = self.name
-        if self.default is not _empty:
+        if self.annotation is not _empty:
+            res += ': ' + self.annotation
+            if self.default is not _empty:
+                res += ' = ' + repr(self.default)
+        elif self.default is not _empty:
             res += '=' + repr(self.default)
         if self.kind == Parameter.VAR_POSITIONAL:
             res = '*' + res
@@ -34,25 +39,34 @@ class Parameter:
 class Signature:
     empty = _empty
 
-    def __init__(self, parameters):
+    def __init__(self, parameters, return_annotation=None):
         self.parameters = {p.name: p for p in parameters}
+        self.return_annotation = _empty if return_annotation is None else return_annotation
 
     def __str__(self):
-        return '(' + ', '.join([str(p) for p in self.parameters.values()]) + ')'
+        res = '(' + ', '.join([str(p) for p in self.parameters.values()]) + ')'
+        if self.return_annotation is not _empty:
+            res += ' -> ' + self.return_annotation
+        return res
 
     def __repr__(self):
         return '<Signature ' + str(self) + '>'
 
 
-def _make_params(func):
-    return [Parameter(*entry) for entry in _signature_data(func)]
+def _from_function(func, drop_first):
+    annotations = func.__annotations__
+    params = [Parameter(*entry, annotation=annotations.get(entry[0]))
+              for entry in _signature_data(func)]
+    if drop_first:
+        params = params[1:]
+    return Signature(params, annotations.get('return'))
 
 
 def signature(obj):
     if not callable(obj):
         raise TypeError(repr(obj) + ' is not a callable object')
     if isinstance(obj, type):
-        return Signature(_make_params(obj.__init__)[1:])
+        return _from_function(obj.__init__, True)
     if hasattr(obj, '__func__'):
-        return Signature(_make_params(obj.__func__)[1:])
-    return Signature(_make_params(obj))
+        return _from_function(obj.__func__, True)
+    return _from_function(obj, False)

src/compiler/compiler.c

@@ -2389,7 +2389,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;
@@ -2438,7 +2442,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

@@ -636,6 +636,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

@@ -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,48 @@ try:
     exit(1)
 except TypeError:
     pass
+
+# ---------------- __annotations__ ----------------
+from inspect import Signature
+
+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',
+}
+
+sig = signature(h1)
+assert sig.parameters['a'].annotation == 'int'
+assert sig.return_annotation == 'bool'
+assert str(sig) == "(a: int, *args: int, b: str = 'x', c: float = 1.5, **kwargs: str) -> bool"
+
+def h2(a, b: int, c=1):
+    pass
+
+assert h2.__annotations__ == {'b': 'int'}
+assert str(signature(h2)) == '(a, b: int, c=1)'
+
+# 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'"}
+
+class D:
+    def m(self, x: int) -> str:
+        return str(x)
+
+sig = signature(D().m)
+assert sig.parameters['x'].annotation == 'int'
+assert sig.return_annotation == 'str'
+
+assert (lambda x: x).__annotations__ == {}
+sig = signature(lambda x: x)
+assert sig.parameters['x'].annotation is Parameter.empty
+assert sig.return_annotation is Signature.empty