PEP 3107 - Function Annotations thanks to Tony Lownds

1 files changed, 53 insertions, 20 deletions

diff --git a/Python/symtable.c b/Python/symtable.c
index b0ebed4..708e18c 100644
--- a/Python/symtable.c
+++ b/Python/symtable.c
@@ -172,9 +172,12 @@ static int symtable_visit_alias(struct symtable *st, alias_ty);
 static int symtable_visit_comprehension(struct symtable *st, comprehension_ty);
 static int symtable_visit_keyword(struct symtable *st, keyword_ty);
 static int symtable_visit_slice(struct symtable *st, slice_ty);
-static int symtable_visit_params(struct symtable *st, asdl_seq *args, int top);
-static int symtable_visit_params_nested(struct symtable *st, asdl_seq *args);
+static int symtable_visit_params(struct symtable *st, asdl_seq *args, int top,
+                                 int annotations);
+static int symtable_visit_params_nested(struct symtable *st, asdl_seq *args,
+                                        int annotations);
 static int symtable_implicit_arg(struct symtable *st, int pos);
+static int symtable_visit_annotations(struct symtable *st, stmt_ty s);
 
 
 static identifier top = NULL, lambda = NULL, genexpr = NULL;
@@ -935,6 +938,8 @@ symtable_visit_stmt(struct symtable *st, stmt_ty s)
 		if (s->v.FunctionDef.args->kw_defaults)
 			VISIT_KWONLYDEFAULTS(st, 
 					   s->v.FunctionDef.args->kw_defaults);
+		if (!symtable_visit_annotations(st, s))
+			return 0;
 		if (s->v.FunctionDef.decorators)
 			VISIT_SEQ(st, expr, s->v.FunctionDef.decorators);
 		if (!symtable_enter_block(st, s->v.FunctionDef.name, 
@@ -1219,22 +1224,29 @@ symtable_implicit_arg(struct symtable *st, int pos)
 }
 
 static int 
-symtable_visit_params(struct symtable *st, asdl_seq *args, int toplevel)
+symtable_visit_params(struct symtable *st, asdl_seq *args, int toplevel,
+                      int annotations)
 {
 	int i;
+
+	if (!args)
+		return -1;
 	
         /* go through all the toplevel arguments first */
 	for (i = 0; i < asdl_seq_LEN(args); i++) {
-		expr_ty arg = (expr_ty)asdl_seq_GET(args, i);
-		if (arg->kind == Name_kind) {
-			assert(arg->v.Name.ctx == Param ||
-                               (arg->v.Name.ctx == Store && !toplevel));
-			if (!symtable_add_def(st, arg->v.Name.id, DEF_PARAM))
-				return 0;
+		arg_ty arg = (arg_ty)asdl_seq_GET(args, i);
+		if (arg->kind == SimpleArg_kind) {
+			if (!annotations) {
+				if (!symtable_add_def(st,
+				                      arg->v.SimpleArg.arg,
+				                      DEF_PARAM))
+					return 0;
+			}
+			else if (arg->v.SimpleArg.annotation)
+				VISIT(st, expr, arg->v.SimpleArg.annotation);
 		}
-		else if (arg->kind == Tuple_kind) {
-			assert(arg->v.Tuple.ctx == Store);
-			if (toplevel) {
+		else if (arg->kind == NestedArgs_kind) {
+			if (toplevel && !annotations) {
 				if (!symtable_implicit_arg(st, i))
 					return 0;
 			}
@@ -1249,7 +1261,7 @@ symtable_visit_params(struct symtable *st, asdl_seq *args, int toplevel)
 	}
 
 	if (!toplevel) {
-		if (!symtable_visit_params_nested(st, args))
+		if (!symtable_visit_params_nested(st, args, annotations))
 			return 0;
 	}
 
@@ -1257,16 +1269,37 @@ symtable_visit_params(struct symtable *st, asdl_seq *args, int toplevel)
 }
 
 static int
-symtable_visit_params_nested(struct symtable *st, asdl_seq *args)
+symtable_visit_params_nested(struct symtable *st, asdl_seq *args,
+                             int annotations)
 {
 	int i;
 	for (i = 0; i < asdl_seq_LEN(args); i++) {
-		expr_ty arg = (expr_ty)asdl_seq_GET(args, i);
-		if (arg->kind == Tuple_kind &&
-		    !symtable_visit_params(st, arg->v.Tuple.elts, 0))
+		arg_ty arg = (arg_ty)asdl_seq_GET(args, i);
+		if (arg->kind == NestedArgs_kind &&
+		    !symtable_visit_params(st, arg->v.NestedArgs.args, 0,
+		                           annotations))
 			return 0;
 	}
+
+	return 1;
+}
+
+
+static int
+symtable_visit_annotations(struct symtable *st, stmt_ty s)
+{
+	arguments_ty a = s->v.FunctionDef.args;
 	
+	if (a->args && !symtable_visit_params(st, a->args, 1, 1))
+		return 0;
+	if (a->varargannotation)
+		VISIT(st, expr, a->varargannotation);
+	if (a->kwargannotation)
+		VISIT(st, expr, a->kwargannotation);
+	if (a->kwonlyargs && !symtable_visit_params(st, a->kwonlyargs, 1, 1))
+		return 0;
+	if (s->v.FunctionDef.returns)
+		VISIT(st, expr, s->v.FunctionDef.returns);
 	return 1;
 }
 
@@ -1276,9 +1309,9 @@ symtable_visit_arguments(struct symtable *st, arguments_ty a)
 	/* skip default arguments inside function block
 	   XXX should ast be different?
 	*/
-	if (a->args && !symtable_visit_params(st, a->args, 1))
+	if (a->args && !symtable_visit_params(st, a->args, 1, 0))
 		return 0;
-	if (a->kwonlyargs && !symtable_visit_params(st, a->kwonlyargs, 1))
+	if (a->kwonlyargs && !symtable_visit_params(st, a->kwonlyargs, 1, 0))
 		return 0;
 	if (a->vararg) {
 		if (!symtable_add_def(st, a->vararg, DEF_PARAM))
@@ -1290,7 +1323,7 @@ symtable_visit_arguments(struct symtable *st, arguments_ty a)
 			return 0;
 		st->st_cur->ste_varkeywords = 1;
 	}
-	if (a->args && !symtable_visit_params_nested(st, a->args))
+	if (a->args && !symtable_visit_params_nested(st, a->args, 0))
 		return 0;
 	return 1;
 }