Ast dump Comprehension structs properly

ncw · ncw · commit a08810b8b992 · 2014-11-23T19:29:01.000Z
diff --git a/ast/dump.go b/ast/dump.go
@@ -8,59 +8,67 @@ import (
 	"github.com/ncw/gpython/py"
 )
 
-// Dump an Ast node as a string
-func Dump(ast Ast) string {
-	if ast == nil {
-		return "<nil>"
-	}
-	name := ast.Type().Name
-	if name == "ExprStmt" {
-		name = "Expr"
+func dumpItem(v interface{}) string {
+	switch x := v.(type) {
+	case py.String:
+		return fmt.Sprintf("'%s'", string(x))
+	case py.Bytes:
+		return fmt.Sprintf("b'%s'", string(x))
+	case Identifier:
+		return fmt.Sprintf("'%s'", string(x))
+	case ModBase:
+	case StmtBase:
+	case ExprBase:
+	case SliceBase:
+	case Pos:
+	case Ast:
+		return Dump(x)
+	case py.I__str__:
+		return string(x.M__str__().(py.String))
+	case Comprehension:
+		return dump(v, "comprehension")
 	}
+	return fmt.Sprintf("%v", v)
+}
+
+// Dump ast as a string with name
+func dump(ast interface{}, name string) string {
 	astValue := reflect.Indirect(reflect.ValueOf(ast))
 	astType := astValue.Type()
 	args := make([]string, 0)
 	for i := 0; i < astType.NumField(); i++ {
 		fieldType := astType.Field(i)
 		fieldValue := astValue.Field(i)
 		fname := strings.ToLower(fieldType.Name)
+		if fname == "stmtbase" || fname == "exprbase" || fname == "modbase" {
+			continue
+		}
 		if fieldValue.Kind() == reflect.Slice && fieldValue.Type().Elem().Kind() != reflect.Uint8 {
 			strs := make([]string, fieldValue.Len())
 			for i := 0; i < fieldValue.Len(); i++ {
 				element := fieldValue.Index(i)
 				if element.CanInterface() {
-					if x, ok := element.Interface().(Ast); ok {
-						strs[i] = Dump(x)
-					} else {
-						strs[i] = fmt.Sprintf("%v", element)
-					}
-				} else {
-					strs[i] = fmt.Sprintf("%v", element)
+					v := element.Interface()
+					strs[i] = dumpItem(v)
 				}
 			}
 			args = append(args, fmt.Sprintf("%s=[%s]", fname, strings.Join(strs, ", ")))
 		} else if fieldValue.CanInterface() {
 			v := fieldValue.Interface()
-			switch x := v.(type) {
-			case py.String:
-				args = append(args, fmt.Sprintf("%s='%s'", fname, string(x)))
-			case py.Bytes:
-				args = append(args, fmt.Sprintf("%s=b'%s'", fname, string(x)))
-			case Identifier:
-				args = append(args, fmt.Sprintf("%s='%s'", fname, string(x)))
-			case ModBase:
-			case StmtBase:
-			case ExprBase:
-			case SliceBase:
-			case Pos:
-			case Ast:
-				args = append(args, fmt.Sprintf("%s=%s", fname, Dump(x)))
-			case py.I__str__:
-				args = append(args, fmt.Sprintf("%s=%s", fname, x.M__str__()))
-			default:
-				args = append(args, fmt.Sprintf("%s=%v", fname, x))
-			}
+			args = append(args, fmt.Sprintf("%s=%s", fname, dumpItem(v)))
 		}
 	}
 	return fmt.Sprintf("%s(%s)", name, strings.Join(args, ", "))
 }
+
+// Dump an Ast node as a string
+func Dump(ast Ast) string {
+	if ast == nil {
+		return "<nil>"
+	}
+	name := ast.Type().Name
+	if name == "ExprStmt" {
+		name = "Expr"
+	}
+	return dump(ast, name)
+}
diff --git a/ast/dump_test.go b/ast/dump_test.go
@@ -23,6 +23,9 @@ func TestDump(t *testing.T) {
 		{&Module{Body: []Stmt{&ExprStmt{Value: &Tuple{}}}}, `Module(body=[Expr(value=Tuple(elts=[], ctx=UnknownExprContext(0)))])`},
 		{&NameConstant{Value: py.True}, `NameConstant(value=True)`},
 		{&Name{Id: Identifier("hello"), Ctx: Load}, `Name(id='hello', ctx=Load())`},
+		{&ListComp{Elt: &Str{S: py.String("potato")}, Generators: []Comprehension{{
+			Target: &Name{Id: Identifier("hello"), Ctx: Load},
+		}}}, `ListComp(elt=Str(s='potato'), generators=[comprehension(target=Name(id='hello', ctx=Load()), iter=<nil>, ifs=[])])`},
 	} {
 		out := Dump(test.in)
 		if out != test.out {
