ast: Walk - fix null pointer deref and increase test coverage

ncw · ncw · commit 611c8dc15f44 · 2015-04-16T22:36:07.000+01:00
diff --git a/ast/walk.go b/ast/walk.go
@@ -73,7 +73,9 @@ func Walk(ast Ast, Visit func(Ast) bool) {
 		// Body          []Stmt
 		// DecoratorList []Expr
 		// Returns       Expr
-		walk(node.Args)
+		if node.Args != nil {
+			walk(node.Args)
+		}
 		walkStmts(node.Body)
 		walkExprs(node.DecoratorList)
 		walk(node.Returns)
@@ -226,7 +228,9 @@ func Walk(ast Ast, Visit func(Ast) bool) {
 	case *Lambda:
 		// Args *Arguments
 		// Body Expr
-		walk(node.Args)
+		if node.Args != nil {
+			walk(node.Args)
+		}
 		walk(node.Body)
 
 	case *IfExp:
@@ -385,18 +389,24 @@ func Walk(ast Ast, Visit func(Ast) bool) {
 		for _, arg := range node.Args {
 			walk(arg)
 		}
-		walk(node.Vararg)
+		if node.Vararg != nil {
+			walk(node.Vararg)
+		}
 		for _, arg := range node.Kwonlyargs {
 			walk(arg)
 		}
 		walkExprs(node.KwDefaults)
-		walk(node.Kwarg)
+		if node.Kwarg != nil {
+			walk(node.Kwarg)
+		}
 		walkExprs(node.Defaults)
 
 	case *Arg:
 		// Arg        Identifier
 		// Annotation Expr
-		walk(node.Annotation)
+		if node.Annotation != nil {
+			walk(node.Annotation)
+		}
 
 	case *Keyword:
 		// Arg   Identifier
diff --git a/ast/walk_test.go b/ast/walk_test.go
@@ -17,7 +17,70 @@ func TestWalk(t *testing.T) {
 		out []string
 	}{
 		{nil, []string{}},
+
+		// An empty one of everything
 		{&Module{}, []string{"*ast.Module"}},
+		{&Interactive{}, []string{"*ast.Interactive"}},
+		{&Expression{}, []string{"*ast.Expression"}},
+		{&Suite{}, []string{"*ast.Suite"}},
+		{&FunctionDef{}, []string{"*ast.FunctionDef"}},
+		{&ClassDef{}, []string{"*ast.ClassDef"}},
+		{&Return{}, []string{"*ast.Return"}},
+		{&Delete{}, []string{"*ast.Delete"}},
+		{&Assign{}, []string{"*ast.Assign"}},
+		{&AugAssign{}, []string{"*ast.AugAssign"}},
+		{&For{}, []string{"*ast.For"}},
+		{&While{}, []string{"*ast.While"}},
+		{&If{}, []string{"*ast.If"}},
+		{&With{}, []string{"*ast.With"}},
+		{&Raise{}, []string{"*ast.Raise"}},
+		{&Try{}, []string{"*ast.Try"}},
+		{&Assert{}, []string{"*ast.Assert"}},
+		{&Import{}, []string{"*ast.Import"}},
+		{&ImportFrom{}, []string{"*ast.ImportFrom"}},
+		{&Global{}, []string{"*ast.Global"}},
+		{&Nonlocal{}, []string{"*ast.Nonlocal"}},
+		{&ExprStmt{}, []string{"*ast.ExprStmt"}},
+		{&Pass{}, []string{"*ast.Pass"}},
+		{&Break{}, []string{"*ast.Break"}},
+		{&Continue{}, []string{"*ast.Continue"}},
+		{&BoolOp{}, []string{"*ast.BoolOp"}},
+		{&BinOp{}, []string{"*ast.BinOp"}},
+		{&UnaryOp{}, []string{"*ast.UnaryOp"}},
+		{&Lambda{}, []string{"*ast.Lambda"}},
+		{&IfExp{}, []string{"*ast.IfExp"}},
+		{&Dict{}, []string{"*ast.Dict"}},
+		{&Set{}, []string{"*ast.Set"}},
+		{&ListComp{}, []string{"*ast.ListComp"}},
+		{&SetComp{}, []string{"*ast.SetComp"}},
+		{&DictComp{}, []string{"*ast.DictComp"}},
+		{&GeneratorExp{}, []string{"*ast.GeneratorExp"}},
+		{&Yield{}, []string{"*ast.Yield"}},
+		{&YieldFrom{}, []string{"*ast.YieldFrom"}},
+		{&Compare{}, []string{"*ast.Compare"}},
+		{&Call{}, []string{"*ast.Call"}},
+		{&Num{}, []string{"*ast.Num"}},
+		{&Str{}, []string{"*ast.Str"}},
+		{&Bytes{}, []string{"*ast.Bytes"}},
+		{&NameConstant{}, []string{"*ast.NameConstant"}},
+		{&Ellipsis{}, []string{"*ast.Ellipsis"}},
+		{&Attribute{}, []string{"*ast.Attribute"}},
+		{&Subscript{}, []string{"*ast.Subscript"}},
+		{&Starred{}, []string{"*ast.Starred"}},
+		{&Name{}, []string{"*ast.Name"}},
+		{&List{}, []string{"*ast.List"}},
+		{&Tuple{}, []string{"*ast.Tuple"}},
+		{&Slice{}, []string{"*ast.Slice"}},
+		{&ExtSlice{}, []string{"*ast.ExtSlice"}},
+		{&Index{}, []string{"*ast.Index"}},
+		{&ExceptHandler{}, []string{"*ast.ExceptHandler"}},
+		{&Arguments{}, []string{"*ast.Arguments"}},
+		{&Arg{}, []string{"*ast.Arg"}},
+		{&Keyword{}, []string{"*ast.Keyword"}},
+		{&Alias{}, []string{"*ast.Alias"}},
+		{&WithItem{}, []string{"*ast.WithItem"}},
+
+		// Excercise the walk* closures
 		{&Module{Body: []Stmt{&Pass{}}}, []string{"*ast.Module", "*ast.Pass"}},
 		{&Module{Body: []Stmt{&Pass{}, &Pass{}}}, []string{"*ast.Module", "*ast.Pass", "*ast.Pass"}},
 		{&Expression{Body: &Num{}}, []string{"*ast.Expression", "*ast.Num"}},
