Clean up handling of function return type annotation.

py/compile.c

@@ -2764,7 +2764,7 @@ void compile_scope(compiler_t *comp, scope_t *scope, pass_kind_t pass) {
             apply_to_single_or_list(comp, pns->nodes[1], PN_typedargslist, compile_scope_func_param);
         }
 
-        assert(MP_PARSE_NODE_IS_NULL(pns->nodes[2])); // 2 is something...
+        // pns->nodes[2] is return/whole function annotation
 
         compile_node(comp, pns->nodes[3]); // 3 is function body
         // emit return if it wasn't the last opcode

py/grammar.h

@@ -33,8 +33,8 @@ DEF_RULE(decorator, nc, and(4), tok(DEL_AT), rule(dotted_name), opt_rule(trailer
 DEF_RULE(decorators, nc, one_or_more, rule(decorator))
 DEF_RULE(decorated, c(decorated), and(2), rule(decorators), rule(decorated_body))
 DEF_RULE(decorated_body, nc, or(2), rule(classdef), rule(funcdef))
-DEF_RULE(funcdef, c(funcdef), and(8), tok(KW_DEF), tok(NAME), tok(DEL_PAREN_OPEN), opt_rule(typedargslist), tok(DEL_PAREN_CLOSE), opt_rule(funcdef_2), tok(DEL_COLON), rule(suite))
-DEF_RULE(funcdef_2, nc, and(2), tok(DEL_MINUS_MORE), rule(test))
+DEF_RULE(funcdef, c(funcdef), and(8), tok(KW_DEF), tok(NAME), tok(DEL_PAREN_OPEN), opt_rule(typedargslist), tok(DEL_PAREN_CLOSE), opt_rule(funcdefrettype), tok(DEL_COLON), rule(suite))
+DEF_RULE(funcdefrettype, nc, and(2), tok(DEL_MINUS_MORE), rule(test))
 // TODO typedargslist lets through more than is allowed
 DEF_RULE(typedargslist, nc, list_with_end, rule(typedargslist_item), tok(DEL_COMMA))
 DEF_RULE(typedargslist_item, nc, or(3), rule(typedargslist_name), rule(typedargslist_star), rule(typedargslist_dbl_star))

tests/basics/fun-annotations.py

@@ -0,0 +1,4 @@
+def foo(x: int, y: list) -> dict:
+    return {x: y}
+
+print(foo(1, [2, 3]))